Tests: refactor trainer (#1728)

* lr

* optim

* wip

* wip

* fix mean

* flake8

tests/base/eval_model_optimizers.py

@@ -12,7 +12,7 @@ class ConfigureOptimizersPool(ABC):
         optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
         return optimizer
 
-    def configure_optimizers_empty(self):
+    def configure_optimizers__empty(self):
         return None
 
     def configure_optimizers__lbfgs(self):

tests/base/eval_model_template.py

@@ -1,3 +1,5 @@
+from argparse import Namespace
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -37,7 +39,7 @@ class EvalModelTemplate(
         """Pass in parsed HyperOptArgumentParser to the model."""
         # init superclass
         super().__init__()
-        self.hparams = hparams
+        self.hparams = Namespace(**hparams) if isinstance(hparams, dict) else hparams
 
         # if you specify an example input, the summary will show input/output for each layer
         self.example_input_array = torch.rand(5, 28 * 28)

tests/base/eval_model_test_dataloaders.py

@@ -9,3 +9,6 @@ class TestDataloaderVariations(ABC):
 
     def test_dataloader(self):
         return self.dataloader(train=False)
+
+    def test_dataloader__empty(self):
+        return None

tests/base/eval_model_train_steps.py

@@ -1,11 +1,16 @@
+import math
 from abc import ABC
 from collections import OrderedDict
 
+import torch
+
 
 class TrainingStepVariations(ABC):
     """
     Houses all variations of training steps
     """
+    test_step_inf_loss = float('inf')
+
     def training_step(self, batch, batch_idx, optimizer_idx=None):
         """Lightning calls this inside the training loop"""
         # forward pass
@@ -28,3 +33,12 @@ class TrainingStepVariations(ABC):
 
         if self.trainer.batch_idx % 2 == 0:
             return loss_val
+
+    def training_step__inf_loss(self, batch, batch_idx, optimizer_idx=None):
+        output = self.training_step(batch, batch_idx, optimizer_idx)
+        if batch_idx == self.test_step_inf_loss:
+            if isinstance(output, dict):
+                output['loss'] *= torch.tensor(math.inf)  # make loss infinite
+            else:
+                output /= 0
+        return output

tests/base/eval_model_valid_dataloaders.py

@@ -9,3 +9,7 @@ class ValDataloaderVariations(ABC):
 
     def val_dataloader(self):
         return self.dataloader(train=False)
+
+    def val_dataloader__multiple(self):
+        return [self.dataloader(train=False),
+                self.dataloader(train=False)]

tests/base/eval_model_valid_epoch_ends.py

@@ -16,9 +16,13 @@ class ValidationEpochEndVariations(ABC):
         """
         # if returned a scalar from validation_step, outputs is a list of tensor scalars
         # we return just the average in this case (if we want)
+        def _mean(res, key):
+            # recursive mean for multilevel dicts
+            return torch.stack([x[key] if isinstance(x, dict) else _mean(x, key) for x in res]).mean()
+
         # return torch.stack(outputs).mean()
-        val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
-        val_acc_mean = torch.stack([x['val_acc'] for x in outputs]).mean()
+        val_loss_mean = _mean(outputs, 'val_loss')
+        val_acc_mean = _mean(outputs, 'val_acc')
         for output in outputs:
             val_loss = self.get_output_metric(output, 'val_loss')
 

tests/base/models.py

@@ -8,6 +8,7 @@ import torch.nn.functional as F
 from torch import optim
 from torch.utils.data import DataLoader
 
+from tests.base import EvalModelTemplate
 from tests.base.datasets import TrialMNIST
 
 try:

tests/trainer/test_lr_finder.py

@@ -4,25 +4,14 @@ import torch
 import tests.base.utils as tutils
 from pytorch_lightning import Trainer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from tests.base import (
-    LightTrainDataloader,
-    TestModelBase,
-    LightTestMultipleOptimizersWithSchedulingMixin,
-)
+from tests.base import EvalModelTemplate
 
 
 def test_error_on_more_than_1_optimizer(tmpdir):
     """ Check that error is thrown when more than 1 optimizer is passed """
 
-    class CurrentTestModel(
-        LightTestMultipleOptimizersWithSchedulingMixin,
-        LightTrainDataloader,
-        TestModelBase,
-    ):
-        pass
-
-    hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
+    model.configure_optimizers = model.configure_optimizers__multiple_schedulers
 
     # logger file to get meta
     trainer = Trainer(
@@ -37,14 +26,7 @@ def test_error_on_more_than_1_optimizer(tmpdir):
 def test_model_reset_correctly(tmpdir):
     """ Check that model weights are correctly reset after lr_find() """
 
-    class CurrentTestModel(
-        LightTrainDataloader,
-        TestModelBase,
-    ):
-        pass
-
-    hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     # logger file to get meta
     trainer = Trainer(
@@ -66,14 +48,7 @@ def test_model_reset_correctly(tmpdir):
 def test_trainer_reset_correctly(tmpdir):
     """ Check that all trainer parameters are reset correctly after lr_find() """
 
-    class CurrentTestModel(
-        LightTrainDataloader,
-        TestModelBase,
-    ):
-        pass
-
-    hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     # logger file to get meta
     trainer = Trainer(
@@ -102,15 +77,10 @@ def test_trainer_reset_correctly(tmpdir):
 
 def test_trainer_arg_bool(tmpdir):
 
-    class CurrentTestModel(
-        LightTrainDataloader,
-        TestModelBase,
-    ):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
     before_lr = hparams.learning_rate
+
     # logger file to get meta
     trainer = Trainer(
         default_save_path=tmpdir,
@@ -126,15 +96,10 @@ def test_trainer_arg_bool(tmpdir):
 
 def test_trainer_arg_str(tmpdir):
 
-    class CurrentTestModel(
-        LightTrainDataloader,
-        TestModelBase,
-    ):
-        pass
-
     hparams = tutils.get_default_hparams()
     hparams.__dict__['my_fancy_lr'] = 1.0  # update with non-standard field
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+
     before_lr = hparams.my_fancy_lr
     # logger file to get meta
     trainer = Trainer(
@@ -151,14 +116,9 @@ def test_trainer_arg_str(tmpdir):
 
 def test_call_to_trainer_method(tmpdir):
 
-    class CurrentTestModel(
-        LightTrainDataloader,
-        TestModelBase,
-    ):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+
     before_lr = hparams.learning_rate
     # logger file to get meta
     trainer = Trainer(

tests/trainer/test_optimizers.py

@@ -3,30 +3,15 @@ import torch
 
 import tests.base.utils as tutils
 from pytorch_lightning import Trainer
-from tests.base import (
-    TestModelBase,
-    LightTrainDataloader,
-    LightValidationStepMixin,
-    LightValidationMixin,
-    LightTestOptimizerWithSchedulingMixin,
-    LightTestMultipleOptimizersWithSchedulingMixin,
-    LightTestOptimizersWithMixedSchedulingMixin,
-    LightTestReduceLROnPlateauMixin,
-    LightTestNoneOptimizerMixin, EvalModelTemplate
-)
+from tests.base import EvalModelTemplate
 
 
 def test_optimizer_with_scheduling(tmpdir):
     """ Verify that learning rate scheduling is working """
 
-    class CurrentTestModel(
-            LightTestOptimizerWithSchedulingMixin,
-            LightTrainDataloader,
-            TestModelBase):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+    model.configure_optimizers = model.configure_optimizers__single_scheduler
 
     # fit model
     trainer = Trainer(
@@ -36,6 +21,7 @@ def test_optimizer_with_scheduling(tmpdir):
         train_percent_check=0.2
     )
     results = trainer.fit(model)
+    assert results == 1
 
     init_lr = hparams.learning_rate
     adjusted_lr = [pg['lr'] for pg in trainer.optimizers[0].param_groups]
@@ -54,14 +40,9 @@ def test_optimizer_with_scheduling(tmpdir):
 def test_multi_optimizer_with_scheduling(tmpdir):
     """ Verify that learning rate scheduling is working """
 
-    class CurrentTestModel(
-            LightTestMultipleOptimizersWithSchedulingMixin,
-            LightTrainDataloader,
-            TestModelBase):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+    model.configure_optimizers = model.configure_optimizers__multiple_schedulers
 
     # fit model
     trainer = Trainer(
@@ -71,6 +52,7 @@ def test_multi_optimizer_with_scheduling(tmpdir):
         train_percent_check=0.2
     )
     results = trainer.fit(model)
+    assert results == 1
 
     init_lr = hparams.learning_rate
     adjusted_lr1 = [pg['lr'] for pg in trainer.optimizers[0].param_groups]
@@ -93,14 +75,9 @@ def test_multi_optimizer_with_scheduling(tmpdir):
 
 def test_multi_optimizer_with_scheduling_stepping(tmpdir):
 
-    class CurrentTestModel(
-            LightTestOptimizersWithMixedSchedulingMixin,
-            LightTrainDataloader,
-            TestModelBase):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+    model.configure_optimizers = model.configure_optimizers__multiple_schedulers
 
     # fit model
     trainer = Trainer(
@@ -110,6 +87,7 @@ def test_multi_optimizer_with_scheduling_stepping(tmpdir):
         train_percent_check=0.2
     )
     results = trainer.fit(model)
+    assert results == 1
 
     init_lr = hparams.learning_rate
     adjusted_lr1 = [pg['lr'] for pg in trainer.optimizers[0].param_groups]
@@ -127,7 +105,7 @@ def test_multi_optimizer_with_scheduling_stepping(tmpdir):
     adjusted_lr2 = adjusted_lr2[0]
 
     # Called ones after end of epoch
-    assert init_lr * 0.1 ** 0 == adjusted_lr1, \
+    assert init_lr * 0.1 ** 1 == adjusted_lr1, \
         'lr for optimizer 1 not adjusted correctly'
     # Called every 3 steps, meaning for 1 epoch of 11 batches, it is called 3 times
     assert init_lr * 0.1 == adjusted_lr2, \
@@ -136,16 +114,9 @@ def test_multi_optimizer_with_scheduling_stepping(tmpdir):
 
 def test_reduce_lr_on_plateau_scheduling(tmpdir):
 
-    class CurrentTestModel(
-            LightTestReduceLROnPlateauMixin,
-            LightTrainDataloader,
-            LightValidationMixin,
-            LightValidationStepMixin,
-            TestModelBase):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+    model.configure_optimizers = model.configure_optimizers__reduce_lr_on_plateau
 
     # fit model
     trainer = Trainer(
@@ -155,7 +126,7 @@ def test_reduce_lr_on_plateau_scheduling(tmpdir):
         train_percent_check=0.2
     )
     results = trainer.fit(model)
-    assert results
+    assert results == 1
 
     assert trainer.lr_schedulers[0] == \
         dict(scheduler=trainer.lr_schedulers[0]['scheduler'], monitor='val_loss',
@@ -233,14 +204,9 @@ def test_none_optimizer_warning():
 
 def test_none_optimizer(tmpdir):
 
-    class CurrentTestModel(
-            LightTestNoneOptimizerMixin,
-            LightTrainDataloader,
-            TestModelBase):
-        pass
-
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(hparams)
+    model.configure_optimizers = model.configure_optimizers__empty
 
     # fit model
     trainer = Trainer(
@@ -256,11 +222,9 @@ def test_none_optimizer(tmpdir):
 
 
 def test_configure_optimizer_from_dict(tmpdir):
-    """Tests if `configure_optimizer` method could return a dictionary with
-    `optimizer` field only.
-    """
+    """Tests if `configure_optimizer` method could return a dictionary with `optimizer` field only."""
 
-    class CurrentTestModel(LightTrainDataloader, TestModelBase):
+    class CurrentModel(EvalModelTemplate):
         def configure_optimizers(self):
             config = {
                 'optimizer': torch.optim.SGD(params=self.parameters(), lr=1e-03)
@@ -268,7 +232,7 @@ def test_configure_optimizer_from_dict(tmpdir):
             return config
 
     hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = CurrentModel(hparams)
 
     # fit model
     trainer = Trainer(default_save_path=tmpdir, max_epochs=1)

tests/trainer/test_trainer.py

@@ -8,34 +8,24 @@ import pytest
 import torch
 
 import tests.base.utils as tutils
-from pytorch_lightning import Callback
+from pytorch_lightning import Callback, LightningModule
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 from pytorch_lightning.core.lightning import load_hparams_from_tags_csv
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from tests.base import (
-    TestModelBase,
-    DictHparamsModel,
-    LightningTestModel,
-    LightEmptyTestStep,
-    LightValidationStepMixin,
-    LightValidationMultipleDataloadersMixin,
-    LightTrainDataloader,
-    LightTestDataloader,
-    LightValidationMixin, EvalModelTemplate,
-)
+from tests.base import EvalModelTemplate
 
 
 def test_model_pickle(tmpdir):
     import pickle
 
-    model = TestModelBase(tutils.get_default_hparams())
+    model = EvalModelTemplate(tutils.get_default_hparams())
     pickle.dumps(model)
 
 
 def test_hparams_save_load(tmpdir):
-    model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10, 'failed_key': lambda x: x})
+    model = EvalModelTemplate(vars(tutils.get_default_hparams()))
 
     trainer = Trainer(
         default_root_dir=tmpdir,
@@ -48,19 +38,15 @@ def test_hparams_save_load(tmpdir):
     # try to load the model now
     pretrained_model = tutils.load_model_from_checkpoint(
         trainer.checkpoint_callback.dirpath,
-        module_class=DictHparamsModel
+        module_class=EvalModelTemplate
     )
+    assert pretrained_model
 
 
 def test_no_val_module(tmpdir):
     """Tests use case where trainer saves the model, and user loads it from tags independently."""
 
-    hparams = tutils.get_default_hparams()
-
-    class CurrentTestModel(LightTrainDataloader, TestModelBase):
-        pass
-
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     # logger file to get meta
     logger = tutils.get_default_logger(tmpdir)
@@ -84,7 +70,7 @@ def test_no_val_module(tmpdir):
     assert 'hparams' in ckpt.keys(), 'hparams missing from checkpoints'
 
     # won't load without hparams in the ckpt
-    model_2 = LightningTestModel.load_from_checkpoint(
+    model_2 = EvalModelTemplate.load_from_checkpoint(
         checkpoint_path=new_weights_path,
     )
     model_2.eval()
@@ -93,11 +79,7 @@ def test_no_val_module(tmpdir):
 def test_no_val_end_module(tmpdir):
     """Tests use case where trainer saves the model, and user loads it from tags independently."""
 
-    class CurrentTestModel(LightTrainDataloader, LightValidationStepMixin, TestModelBase):
-        pass
-
-    hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     # logger file to get meta
     logger = tutils.get_default_logger(tmpdir)
@@ -120,7 +102,7 @@ def test_no_val_end_module(tmpdir):
     # load new model
     tags_path = tutils.get_data_path(logger, path_dir=tmpdir)
     tags_path = os.path.join(tags_path, 'meta_tags.csv')
-    model_2 = LightningTestModel.load_from_checkpoint(
+    model_2 = EvalModelTemplate.load_from_checkpoint(
         checkpoint_path=new_weights_path,
         tags_csv=tags_path
     )
@@ -185,8 +167,7 @@ def test_gradient_accumulation_scheduling(tmpdir):
         # clear gradients
         optimizer.zero_grad()
 
-    hparams = tutils.get_default_hparams()
-    model = LightningTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
     schedule = {1: 2, 3: 4}
 
     trainer = Trainer(accumulate_grad_batches=schedule,
@@ -260,9 +241,6 @@ def test_model_checkpoint_options(tmpdir, save_top_k, file_prefix, expected_file
     def mock_save_function(filepath):
         open(filepath, 'a').close()
 
-    hparams = tutils.get_default_hparams()
-    _ = LightningTestModel(hparams)
-
     # simulated losses
     losses = [10, 9, 2.8, 5, 2.5]
 
@@ -288,8 +266,7 @@ def test_model_checkpoint_options(tmpdir, save_top_k, file_prefix, expected_file
 
 def test_model_freeze_unfreeze():
 
-    hparams = tutils.get_default_hparams()
-    model = LightningTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     model.freeze()
     model.unfreeze()
@@ -302,7 +279,7 @@ def test_resume_from_checkpoint_epoch_restored(tmpdir):
 
     def _new_model():
         # Create a model that tracks epochs and batches seen
-        model = LightningTestModel(hparams)
+        model = EvalModelTemplate(hparams)
         model.num_epochs_seen = 0
         model.num_batches_seen = 0
         model.num_on_load_checkpoint_called = 0
@@ -452,15 +429,8 @@ def test_trainer_min_steps_and_epochs(tmpdir):
 def test_benchmark_option(tmpdir):
     """Verify benchmark option."""
 
-    class CurrentTestModel(
-        LightValidationMultipleDataloadersMixin,
-        LightTrainDataloader,
-        TestModelBase
-    ):
-        pass
-
-    hparams = tutils.get_default_hparams()
-    model = CurrentTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
+    model.val_dataloader = model.val_dataloader__multiple
 
     # verify torch.backends.cudnn.benchmark is not turned on
     assert not torch.backends.cudnn.benchmark
@@ -481,40 +451,34 @@ def test_benchmark_option(tmpdir):
 
 
 def test_testpass_overrides(tmpdir):
+    # todo: check duplicated tests against trainer_checks
     hparams = tutils.get_default_hparams()
 
-    class LocalModel(LightTrainDataloader, TestModelBase):
-        pass
-
-    class LocalModelNoEnd(LightTrainDataloader, LightTestDataloader, LightEmptyTestStep, TestModelBase):
-        pass
-
-    class LocalModelNoStep(LightTrainDataloader, TestModelBase):
-        def test_epoch_end(self, outputs):
-            return {}
-
     # Misconfig when neither test_step or test_end is implemented
-    with pytest.raises(MisconfigurationException):
-        model = LocalModel(hparams)
+    with pytest.raises(MisconfigurationException, match='.*not implement `test_dataloader`.*'):
+        model = EvalModelTemplate(hparams)
+        model.test_dataloader = model.test_dataloader__empty
         Trainer().test(model)
 
     # Misconfig when neither test_step or test_end is implemented
     with pytest.raises(MisconfigurationException):
-        model = LocalModelNoStep(hparams)
+        model = EvalModelTemplate(hparams)
+        model.test_step = LightningModule.test_step
         Trainer().test(model)
 
     # No exceptions when one or both of test_step or test_end are implemented
-    model = LocalModelNoEnd(hparams)
+    model = EvalModelTemplate(hparams)
+    model.test_step_end = LightningModule.test_step_end
     Trainer().test(model)
 
-    model = LightningTestModel(hparams)
+    model = EvalModelTemplate(hparams)
     Trainer().test(model)
 
 
 def test_disabled_validation():
     """Verify that `val_percent_check=0` disables the validation loop unless `fast_dev_run=True`."""
 
-    class CurrentModel(LightTrainDataloader, LightValidationMixin, TestModelBase):
+    class CurrentModel(EvalModelTemplate):
 
         validation_step_invoked = False
         validation_epoch_end_invoked = False
@@ -564,59 +528,56 @@ def test_disabled_validation():
 
 
 def test_nan_loss_detection(tmpdir):
-    test_step = 8
 
-    class InfLossModel(LightTrainDataloader, TestModelBase):
+    class CurrentModel(EvalModelTemplate):
+        test_batch_inf_loss = 8
 
-        def training_step(self, batch, batch_idx):
-            output = super().training_step(batch, batch_idx)
-            if batch_idx == test_step:
+        def training_step(self, batch, batch_idx, optimizer_idx=None):
+            output = super().training_step(batch, batch_idx, optimizer_idx)
+            if batch_idx == self.test_batch_inf_loss:
                 if isinstance(output, dict):
                     output['loss'] *= torch.tensor(math.inf)  # make loss infinite
                 else:
                     output /= 0
             return output
 
-    hparams = tutils.get_default_hparams()
-    model = InfLossModel(hparams)
+    model = CurrentModel(tutils.get_default_hparams())
 
     # fit model
     trainer = Trainer(
         default_root_dir=tmpdir,
-        max_steps=(test_step + 1),
+        max_steps=(model.test_batch_inf_loss + 1),
         terminate_on_nan=True
     )
 
     with pytest.raises(ValueError, match=r'.*The loss returned in `training_step` is nan or inf.*'):
         trainer.fit(model)
-        assert trainer.global_step == test_step
+        assert trainer.global_step == model.test_step_inf_loss
 
     for param in model.parameters():
         assert torch.isfinite(param).all()
 
 
 def test_nan_params_detection(tmpdir):
-    test_step = 8
 
-    class NanParamModel(LightTrainDataloader, TestModelBase):
+    class CurrentModel(EvalModelTemplate):
+        test_batch_nan = 8
 
         def on_after_backward(self):
-            if self.global_step == test_step:
+            if self.global_step == self.test_batch_nan:
                 # simulate parameter that became nan
                 torch.nn.init.constant_(self.c_d1.bias, math.nan)
 
-    hparams = tutils.get_default_hparams()
-
-    model = NanParamModel(hparams)
+    model = CurrentModel(tutils.get_default_hparams())
     trainer = Trainer(
         default_root_dir=tmpdir,
-        max_steps=(test_step + 1),
+        max_steps=(model.test_batch_nan + 1),
         terminate_on_nan=True
     )
 
     with pytest.raises(ValueError, match=r'.*Detected nan and/or inf values in `c_d1.bias`.*'):
         trainer.fit(model)
-        assert trainer.global_step == test_step
+        assert trainer.global_step == model.test_batch_nan
 
     # after aborting the training loop, model still has nan-valued params
     params = torch.cat([param.view(-1) for param in model.parameters()])
@@ -626,7 +587,7 @@ def test_nan_params_detection(tmpdir):
 def test_trainer_interrupted_flag(tmpdir):
     """Test the flag denoting that a user interrupted training."""
 
-    model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10})
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     class InterruptCallback(Callback):
         def __init__(self):
@@ -656,8 +617,7 @@ def test_gradient_clipping(tmpdir):
     Test gradient clipping
     """
 
-    hparams = tutils.get_default_hparams()
-    model = LightningTestModel(hparams)
+    model = EvalModelTemplate(tutils.get_default_hparams())
 
     # test that gradient is clipped correctly
     def _optimizer_step(*args, **kwargs):