add set_epoch for MpDeviceLoaderWrapper (#3053)

* add set_epoch for MpDeviceLoaderWrapper * fix one over-indented space
huggingface · Sep 2, 2024 · 758d624 · 758d624
1 parent b07ad2a
commit 758d624
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/accelerate/data_loader.py b/src/accelerate/data_loader.py
@@ -652,6 +652,10 @@ def __iter__(self):
 
             return super().__iter__()
 
+        def set_epoch(self, epoch: int):
+            if hasattr(self.dataloader, "set_epoch"):
+                self.dataloader.set_epoch(epoch)
+
         @property
         def total_batch_size(self):
             return self._loader.total_batch_size