diff --git a/tests/test_models.py b/tests/test_models.py index c958b0d0..4195f9a6 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -36,9 +36,6 @@ def test_cpu_model(): model, hparams = get_model() - # test memory gathering - model.count_mem_items() - run_gpu_model_test(trainer_options, model, hparams, on_gpu=False) @@ -237,6 +234,9 @@ def test_ddp_sampler_error(): use_amp=True ) + # test memory gathering + trainer.count_mem_items() + with pytest.raises(MisconfigurationException): trainer.get_dataloaders(model)