mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-10-27 13:12:20 +08:00
added single gpu train test
This commit is contained in:
1 parent
e30514922b
commit
9ecb1f2aee
1 file changed
+5
@@ -392,6 +392,9 @@ class Trainer(TrainerIO):
|
||||
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 @@ We recommend you switch to ddp if you want to use amp
|
||||
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)
|
||||
|
||||
|
||||
Reference in new issue
Block a user