wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-10 12:21:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

added single gpu train test

Browse Source
This commit is contained in:
William Falcon
2019-08-07 13:50:12 -04:00
parent 2f7a9ad40d
commit 56f16694c4
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/models/trainer.py
+2 -2
View File
@@ -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)