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:54:32 -04:00
parent ab499573a1
commit f3dea818f2
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/models/trainer.py
+1 -1
View File
@@ -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)