From 3521e872868c2d2cdc381f15b7501523a932ea60 Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 24 Jul 2019 17:18:58 -0400
Subject: [PATCH] added multiple outputs to LightningTestModel

---
 pytorch_lightning/models/trainer.py           |  2 +-
 .../testing_models/lm_test_module.py          | 23 ++++++++++++++-----
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
index 4f4d012d..0cd4d89a 100644
--- a/pytorch_lightning/models/trainer.py
+++ b/pytorch_lightning/models/trainer.py
@@ -225,7 +225,7 @@ class Trainer(TrainerIO):
         if self.use_amp:
             print('using 16bit precision')
 
-        if use_amp and not APEX_AVAILABLE:
+        if use_amp and not APEX_AVAILABLE:  # pragma: no cover
             msg = '''
             You set use_amp=True but do not have apex installed.
             Install apex first using this guide and rerun with use_amp=True: 
diff --git a/pytorch_lightning/testing_models/lm_test_module.py b/pytorch_lightning/testing_models/lm_test_module.py
index 3402f4ee..c0843667 100644
--- a/pytorch_lightning/testing_models/lm_test_module.py
+++ b/pytorch_lightning/testing_models/lm_test_module.py
@@ -126,13 +126,23 @@ class LightningTestModel(LightningModule):
             loss_val = loss_val.unsqueeze(0)
             val_acc = val_acc.unsqueeze(0)
 
-        output = OrderedDict({
-            'val_loss': loss_val,
-            'val_acc': val_acc,
-        })
+        # alternate possible outputs to test
+        if self.trainer.batch_nb % 0 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if self.trainer.batch_nb % 1 == 0:
+            return val_acc
 
-        # can also return just a scalar instead of a dict (return loss_val)
-        return output
+        if self.trainer.batch_nb % 2 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
 
     def validation_end(self, outputs):
         """
@@ -152,6 +162,7 @@ class LightningTestModel(LightningModule):
 
         val_loss_mean /= len(outputs)
         val_acc_mean /= len(outputs)
+
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic