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

updated args

Browse Source
This commit is contained in:
William Falcon
2019-06-25 19:46:49 -04:00
parent 158aca26e2
commit 0795e4d51b
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/pt_overrides/override_data_parallel.py
+4 -2
View File
@@ -9,7 +9,6 @@ class LightningDataParallel(DataParallel):
"""
def forward(self, *inputs, **kwargs):
pdb.set_trace()
if not self.device_ids:
# -------------
# MAIN CHANGE
@@ -27,7 +26,10 @@ class LightningDataParallel(DataParallel):
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
if len(self.device_ids) == 1:
return self.module(*inputs[0], **kwargs[0])
if self.module.training:
return self.module.training_step(*inputs[0], **kwargs[0])
else:
return self.module.validation_step(*inputs[0], **kwargs[0])
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
outputs = self.parallel_apply(replicas, inputs, kwargs)
return self.gather(outputs, self.output_device)