From da842c0cd6509ff3aa3bb0eb9b0fdd94b8e0c7b5 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Thu, 18 Jul 2019 12:04:45 -0400
Subject: [PATCH] added arg docs

---
 examples/new_project_templates/lightning_module_template.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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):
         """