Fix issue with tests (#2111)

huggingface · Nov 1, 2023 · c2d8e24 · c2d8e24
1 parent d8e1285
commit c2d8e24
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/test_utils/scripts/test_script.py b/src/accelerate/test_utils/scripts/test_script.py
@@ -339,7 +339,7 @@ def mock_training(length, batch_size, generator):
     set_seed(42)
     generator.manual_seed(42)
     train_set = RegressionDataset(length=length, seed=42)
-    if AcceleratorState().num_processes > 1:
+    if not ((AcceleratorState().num_processes == 1) and (AcceleratorState().device.type == "cpu")):
         # The SeedableRandomSampler is needed during distributed setups
         # for full reproducability across processes with the `DataLoader`
         sampler = SeedableRandomSampler(