Fix finite IterableDataset test on multiple GPUs (#14445)

huggingface · Nov 18, 2021 · 83ef8bc · 83ef8bc
1 parent da36c55
commit 83ef8bc
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/test_trainer.py b/tests/test_trainer.py
@@ -1069,13 +1069,17 @@ def test_training_iterable_dataset(self):
         self.assertIsInstance(loader.sampler, torch.utils.data.dataloader._InfiniteConstantSampler)
 
     def test_training_finite_iterable_dataset(self):
+        num_gpus = max(1, get_gpu_count())
+        if num_gpus > 2:
+            return
+
         config = RegressionModelConfig()
         model = RegressionPreTrainedModel(config)
 
         batch_size = 1
         num_samples = 10
 
-        available_steps = num_samples // batch_size
+        available_steps = num_samples // (batch_size * num_gpus)
 
         data = FiniteIterableDataset(length=num_samples)
         train_args = TrainingArguments(