pickling generator issues

byi8220 · Sep 4, 2024 · 40ec962 · 40ec962
1 parent b43c85d
commit 40ec962
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/accelerate/test_utils/scripts/test_distributed_data_loop.py b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
@@ -74,7 +74,7 @@ def __iter__(self):
 def create_accelerator(even_batches=True):
     dataloader_config = DataLoaderConfiguration(even_batches=even_batches)
     accelerator = Accelerator(dataloader_config=dataloader_config)
-    # assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
+    assert accelerator.num_processes == 2, "this script expects that two GPUs are available"
     return accelerator
 
 
@@ -341,8 +341,6 @@ def test_stateful_dataloader_save_state(accelerator):
         accelerator.dataloader_config = old_dataloader_config
 
 def test_pickled_dataloader(data_loader, accelerator):
-    # Prepare the DataLoader
-    data_loader = accelerator.prepare(data_loader)
     # Pickle then reload the dataloader
     prepared_model_dumps = pickle.dumps(accelerator)
     loaded_accelerator = pickle.loads(prepared_model_dumps)

diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
@@ -662,7 +662,7 @@ def test_can_pickle_dataloader(self, dispatch_batches):
         # TODO: Add support for pickling StatefulDataLoader
         dataloader_config = DataLoaderConfiguration(dispatch_batches=dispatch_batches, use_stateful_dataloader=False)
         accelerator = Accelerator(dataloader_config=dataloader_config)
-        torch.manual_seed(accelerator.process_index)
+
         original_dl, prepared_skip_dl = accelerator.prepare(dl, skip_dl)
         prepared_model_dumps = pickle.dumps(accelerator)