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

refactored model tests

Browse Source
This commit is contained in:
William Falcon
2019-07-24 14:04:17 -04:00
parent d004fc5725
commit 42b86a160d
1 changed files with 12 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/examples/new_project_templates/lightning_module_template.py
+12 -3
View File
@@ -84,8 +84,12 @@ class LightningTemplateModel(LightningModule):
# calculate loss
loss_val = self.loss(y, y_hat)
# in DP mode (default) make sure if result is scalar, there's another dim in the beginning
if self.trainer.use_dp:
loss_val = loss_val.unsqueeze(0)
output = OrderedDict({
'loss': loss_val.unsqueeze(0)
'loss': loss_val
})
# can also return just a scalar instead of a dict (return loss_val)
@@ -111,9 +115,14 @@ class LightningTemplateModel(LightningModule):
if self.on_gpu:
val_acc = val_acc.cuda(loss_val.device.index)
# in DP mode (default) make sure if result is scalar, there's another dim in the beginning
if self.trainer.use_dp:
loss_val = loss_val.unsqueeze(0)
val_acc = val_acc.unsqueeze(0)
output = OrderedDict({
'val_loss': loss_val.unsqueeze(0),
'val_acc': val_acc.unsqueeze(0),
'val_loss': loss_val,
'val_acc': val_acc,
})
# can also return just a scalar instead of a dict (return loss_val)