added clean slurm save load test

2026-09-11 12:31:23 +08:00 · 2019-07-26 22:33:00 -04:00
parent f5a01edfb8
commit 8e3a0443c7
1 changed files with 1 additions and 0 deletions
@@ -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