added clean slurm save load test

2019-07-26 22:40:07 -04:00 · 2019-07-26 22:40:07 -04:00 · 61c82611eb
parent 3224365190
commit 61c82611eb
1 changed files with 1 additions and 2 deletions
--- a/tests/test_models.py
+++ b/tests/test_models.py
@ -80,9 +80,8 @@ def test_cpu_slurm_save_load():
    trainer.model = LightningTestModel(hparams)

    # test HPC loading
-    trainer.global_step = 20000000
    trainer.hpc_load(save_dir, on_gpu=False)
-    assert trainer.global_step == real_global_step and trainer.global_step != 20000000
+    assert trainer.global_step == real_global_step and trainer.global_step > 0

    # predict with loaded model to make sure answers are the same
    trainer.model.eval()