From 56f16694c4e85894e5e0d8e87c87dbf1112a86ab Mon Sep 17 00:00:00 2001 From: William Falcon Date: Wed, 7 Aug 2019 13:50:12 -0400 Subject: [PATCH] added single gpu train test --- pytorch_lightning/models/trainer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 91eb7cc8..37fd9d8a 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -395,7 +395,7 @@ class Trainer(TrainerIO): elif self.single_gpu: gpu_id = self.data_parallel_device_ids[0] data_batch = [x.cuda(gpu_id) for x in data_batch if isinstance(x, torch.Tensor)] - output = model(data_batch, batch_i) + output = model.validation_step(data_batch, batch_i) else: output = model.validation_step(data_batch, batch_i) @@ -850,7 +850,7 @@ We recommend you switch to ddp if you want to use amp elif self.single_gpu: gpu_id = self.data_parallel_device_ids[0] data_batch = [x.cuda(gpu_id) for x in data_batch if isinstance(x, torch.Tensor)] - output = self.model(data_batch.cuda(self.data_parallel_device_ids[0]), batch_nb) + output = self.model.training_step(data_batch.cuda(self.data_parallel_device_ids[0]), batch_nb) else: output = self.model.training_step(data_batch, batch_nb)