added single gpu train test

Author: williamFalcon

pytorch_lightning/models/trainer.py

@@ -392,6 +392,9 @@ def validate(self, model, dataloader, max_batches):
                 output = model(data_batch, batch_i)
                 output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
 
+            elif self.single_gpu:
+                output = model(data_batch.cuda(self.data_parallel_device_ids[0]), batch_i)
+
             else:
                 output = model.validation_step(data_batch, batch_i)
 
@@ -842,6 +845,8 @@ def __run_tng_batch(self, data_batch, batch_nb):
         elif self.use_dp:
             output = self.model(data_batch, batch_nb)
             output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+        elif self.single_gpu:
+            output = self.model(data_batch.cuda(self.data_parallel_device_ids[0]), batch_nb)
         else:
             output = self.model.training_step(data_batch, batch_nb)