SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Fix finite IterableDataset test on multiple GPUs (#14445)

Browse Source
This commit is contained in:
Sylvain Gugger
2021-11-18 10:25:06 -05:00
committed by GitHub
parent da36c557f7
commit 83ef8bcac2
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_trainer.py
View File

@@ -1069,13 +1069,17 @@ class TrainerIntegrationTest(TestCasePlus, TrainerIntegrationCommon):
self.assertIsInstance(loader.sampler, torch.utils.data.dataloader._InfiniteConstantSampler) self.assertIsInstance(loader.sampler, torch.utils.data.dataloader._InfiniteConstantSampler)
def test_training_finite_iterable_dataset(self): def test_training_finite_iterable_dataset(self):
num_gpus = max(1, get_gpu_count())
if num_gpus > 2:
return
config = RegressionModelConfig() config = RegressionModelConfig()
model = RegressionPreTrainedModel(config) model = RegressionPreTrainedModel(config)
batch_size = 1 batch_size = 1
num_samples = 10 num_samples = 10
available_steps = num_samples // batch_size available_steps = num_samples // (batch_size * num_gpus)
data = FiniteIterableDataset(length=num_samples) data = FiniteIterableDataset(length=num_samples)
train_args = TrainingArguments( train_args = TrainingArguments(