Fix dataloader skipping other batches

huggingface · May 22, 2023 · 6d5d7b0 · 6d5d7b0
1 parent bfa74e5
commit 6d5d7b0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
@@ -369,6 +369,7 @@ def __iter__(self):
                 self.gradient_state._remove_dataloader(self)
                 if batch_index >= self.skip_batches:
                     yield current_batch
+                self.skip_batches = 0
                 break
 
     @property
@@ -565,6 +566,7 @@ def __iter__(self):
             if stop_iteration:
                 self.gradient_state._remove_dataloader(self)
                 self.gradient_state._set_remainder(observed_batch_size)
+                self.skip_batches = 0
             if batch_index >= self.skip_batches:
                 yield batch
             batch_index += 1
@@ -792,6 +794,7 @@ def __iter__(self):
         for index, samples in enumerate(self.batch_sampler):
             if index >= self.skip_batches:
                 yield samples
+        self.skip_batches = 0
 
     @property
     def total_length(self):