From 7a7a9a9da0f57a0302cd8b3e9d56634c6eafd941 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Tue, 25 Jun 2019 20:14:29 -0400
Subject: [PATCH] updated args

---
 docs/source/examples/example_model.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/source/examples/example_model.py b/docs/source/examples/example_model.py
index a8954d90..416344a9 100644
--- a/docs/source/examples/example_model.py
+++ b/docs/source/examples/example_model.py
@@ -76,7 +76,7 @@ class ExampleModel(RootModule):
         # return loss_val, tqdm_dic
 
         output = OrderedDict({
-            'loss_val': loss_val,
+            'loss_val': loss_val.unsqueeze(0),
         })
         return output
 
@@ -99,8 +99,8 @@ class ExampleModel(RootModule):
 
         # output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
         output = OrderedDict({
-            'loss_val': loss_val,
-            'val_acc': val_acc,
+            'loss_val': loss_val.unsqueeze(0),
+            'val_acc': val_acc.unsqueeze(0),
         })
         return output