From f3dea818f27c764fb2fb6c1f27aeed2ea99b4a79 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 7 Aug 2019 13:54:32 -0400 Subject: [PATCH] added single gpu train test --- pytorch_lightning/models/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 55f013b6..40767410 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -854,7 +854,7 @@ We recommend you switch to ddp if you want to use amp for i, x in enumerate(data_batch): if isinstance(x, torch.Tensor): data_batch[i] = x.cuda(gpu_id) - output = self.model.training_step(data_batch.cuda(self.data_parallel_device_ids[0]), batch_nb) + output = self.model.training_step(data_batch, batch_nb) else: output = self.model.training_step(data_batch, batch_nb)