Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fixed reference

This commit is contained in:
SeanNaren 2020-11-22 13:59:02 +00:00
parent 1e429bae58
commit 4ae6f0969a
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
pytorch_lightning/overrides/fairscale.py
View File

@ -21,10 +21,10 @@ class LightningShardedDataParallel(ShardedDataParallel):
if self.enable_broadcast_buffers:
self.sync_buffers()
if self.base_model.training:
outputs = self.base_model.training_step(*inputs, **kwargs)
elif self.base_model.testing:
outputs = self.base_model.test_step(*inputs, **kwargs)
if self.module.training:
outputs = self.module.training_step(*inputs, **kwargs)
elif self.module.testing:
outputs = self.module.test_step(*inputs, **kwargs)
else:
outputs = self.base_model.validation_step(*inputs, **kwargs)
outputs = self.module.validation_step(*inputs, **kwargs)
return outputs