Clean

huggingface · Oct 30, 2023 · 631e824 · 631e824
1 parent 2ab2b1e
commit 631e824
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
@@ -328,6 +328,7 @@ def __len__(self):
     assert isinstance(
         dl.batch_sampler, CustomBatchSampler
     ), "Custom sampler was changed after calling `prepare_data_loader`"
+    state._reset_state()
 
 
 def mock_training(length, batch_size, generator):