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

---
 tests/test_models.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/test_models.py b/tests/test_models.py
index 553a5ccc0a..67c677a248 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -190,7 +190,9 @@ def test_amp_gpu_ddp():
 
     os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
 
-    model, hparams = get_model()
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
     trainer_options = dict(
         progress_bar=True,
         max_nb_epochs=1,