From 16e50938055c73a1f3d92344e1dfa28f03d6134e Mon Sep 17 00:00:00 2001
From: William Falcon <falc@fb.com>
Date: Wed, 7 Aug 2019 11:47:05 -0400
Subject: [PATCH] added test model to do also

---
 pytorch_lightning/testing/lm_test_module.py | 2 +-
 tests/test_models.py                        | 8 ++++++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/pytorch_lightning/testing/lm_test_module.py b/pytorch_lightning/testing/lm_test_module.py
index 8fe4cfc0..695b495f 100644
--- a/pytorch_lightning/testing/lm_test_module.py
+++ b/pytorch_lightning/testing/lm_test_module.py
@@ -232,7 +232,7 @@ class LightningTestModel(LightningModule):
         return self.__dataloader(train=False)
 
     @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):
+    def add_model_specific_args(parent_parser, root_dir):  # pragma: no cover
         """
         Parameters you define here will be available to your model through self.hparams
         :param parent_parser:
diff --git a/tests/test_models.py b/tests/test_models.py
index f6e792fe..ce1697a9 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -681,10 +681,14 @@ def get_hparams(continue_training=False, hpc_exp_number=0):
     return hparams
 
 
-def get_model():
+def get_model(use_test_model=False):
     # set up model with these hyperparams
     hparams = get_hparams()
-    model = LightningTemplateModel(hparams)
+
+    if use_test_model:
+        model = LightningTestModel(hparams)
+    else:
+        model = LightningTemplateModel(hparams)
 
     return model, hparams