diff --git a/examples/new_project_templates/lightning_module_template.py b/examples/new_project_templates/lightning_module_template.py index 2371ed92..c3ecaff5 100644 --- a/examples/new_project_templates/lightning_module_template.py +++ b/examples/new_project_templates/lightning_module_template.py @@ -86,7 +86,7 @@ class LightningTemplateModel(LightningModule): output = OrderedDict({ 'loss': loss_val }) - return output + return loss_val def validation_step(self, data_batch, batch_i): """ @@ -108,7 +108,7 @@ class LightningTemplateModel(LightningModule): 'val_loss': loss_val, 'val_acc': torch.tensor(val_acc).cuda(loss_val.device.index), }) - return output + return loss_val def validation_end(self, outputs): """