fix typo

byi8220 · Sep 3, 2024 · dff2666 · dff2666
1 parent 3235e99
commit dff2666
Showing 1 changed file with 2 additions 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,8 @@ def main():
     test_stateful_dataloader_save_state(accelerator)
 
     # Dataloader after pickling
-    loader = DataLoader(DummyIterableDataset(), shuffle=False, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
+    iterable_loader = create_dataloader(accelerator, dataset_size=NUM_ELEMENTS, batch_size=BATCH_SIZE, iterable=True)
+    loader = DataLoader(DummyIterableDataset(iterable_loader), shuffle=False, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
     test_pickled_dataloader(loader, accelerator)
 
     accelerator.end_training()