Update trainer.py

chaoyi-wu · Mar 25, 2024 · e8afb76 · e8afb76
1 parent 5505f3c
commit e8afb76
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/My_Trainer/trainer.py b/src/My_Trainer/trainer.py
@@ -1880,9 +1880,10 @@ def _inner_training_loop(
 
         total_batched_samples = 0
         for epoch in range(epochs_trained, num_train_epochs):
-            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
+            ### 吴超逸加 ###
+            if isinstance(train_dataloader, DataLoader) and (isinstance(train_dataloader.sampler, DistributedSampler) or self.args.data_sampler != None):
                 train_dataloader.sampler.set_epoch(epoch)
-            elif hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDatasetShard):
+            elif hasattr(train_dataloader, "dataset") and (isinstance(train_dataloader.sampler, DistributedSampler) or self.args.data_sampler != None):
                 train_dataloader.dataset.set_epoch(epoch)
 
             if is_torch_tpu_available():