Fix multiGPU Trainer test by indexing iterator
@@ -79,8 +79,10 @@ class TestTrainer(AllenNlpTestCase):
@pytest.mark.skipif(torch.cuda.device_count() < 2,
reason="Need multiple GPUs.")
def test_trainer_can_run_multiple_gpu(self):
+ multigpu_iterator = BasicIterator(batch_size=4)
+ multigpu_iterator.index_with(self.vocab)
trainer = Trainer(self.model, self.optimizer,
- BasicIterator(batch_size=4), self.instances, num_epochs=2,
+ multigpu_iterator, self.instances, num_epochs=2,
cuda_device=[0, 1])
trainer.train()
