added clean slurm save load test

2019-07-26 22:33:00 -04:00 · 2019-07-26 22:33:00 -04:00 · 8e3a0443c7
parent f5a01edfb8
commit 8e3a0443c7
1 changed files with 1 additions and 0 deletions
--- a/tests/test_models.py
+++ b/tests/test_models.py
@ -182,6 +182,7 @@ def test_cpu_slurm_save_load():
    assert trainer.global_step == real_global_step and trainer.global_step != 20000000

    # predict with loaded model to make sure answers are the same
+    trainer.model.eval()
    new_pred = trainer.model(x)
    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1