fix tests (Lightning-AI#938)

* fix tests * fix tests
tullie · Apr 3, 2020 · 51e2fb3 · 51e2fb3
1 parent 8b25d8c
commit 51e2fb3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 12 deletions.
diff --git a/pytorch_lightning/trainer/evaluation_loop.py b/pytorch_lightning/trainer/evaluation_loop.py
@@ -301,8 +301,8 @@ def evaluate(self, model, dataloaders, max_batches, test=False):
 
     def run_evaluation(self, test=False):
         # when testing make sure user defined a test step
-        if test and not (self.is_overriden('test_step') or self.is_overriden('test_end')):
-            m = '''You called `.test()` without defining model's `.test_step()` or `.test_end()`.
+        if test and not self.is_overriden('test_step'):
+            m = '''You called `.test()` without defining model's `.test_step()`.
                     Please define and try again'''
             raise MisconfigurationException(m)
 

diff --git a/tests/test_trainer.py b/tests/test_trainer.py
@@ -793,22 +793,12 @@ def test_step(self, *args, **kwargs):
         def test_dataloader(self):
             return self.train_dataloader()
 
-    class TestModelNoStep(LightningTestModelBase):
-        def test_end(self, outputs):
-            return {}
-
-        def test_dataloader(self):
-            return self.train_dataloader()
-
     # Misconfig when neither test_step or test_end is implemented
     with pytest.raises(MisconfigurationException):
         model = LightningTestModelBase(hparams)
         Trainer().test(model)
 
     # No exceptions when one or both of test_step or test_end are implemented
-    model = TestModelNoStep(hparams)
-    Trainer().test(model)
-
     model = TestModelNoEnd(hparams)
     Trainer().test(model)