does this work with iterable dsets?

byi8220 · Sep 3, 2024 · 3235e99 · 3235e99
1 parent 99f2afe
commit 3235e99
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/test_utils/scripts/test_distributed_data_loop.py b/src/accelerate/test_utils/scripts/test_distributed_data_loop.py
@@ -411,7 +411,7 @@ def main():
     test_stateful_dataloader_save_state(accelerator)
 
     # Dataloader after pickling
-    loader = DataLoader(dataset, shuffle=False, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
+    loader = DataLoader(DummyIterableDataset(), shuffle=False, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
     test_pickled_dataloader(loader, accelerator)
 
     accelerator.end_training()