use best model dir directly in tests

apax-hub · Nov 10, 2023 · a2c92c9 · a2c92c9
1 parent db80629
commit a2c92c9
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/tests/integration_tests/md/test_md.py b/tests/integration_tests/md/test_md.py
@@ -80,11 +80,8 @@ def test_run_md(get_tmp_path):
     )
 
     ckpt = {"model": {"params": params}, "epoch": 0}
-    best_dir = os.path.join(
-        model_config.data.directory, model_config.data.experiment, "best"
-    )
     checkpoints.save_checkpoint(
-        ckpt_dir=best_dir,
+        ckpt_dir=model_config.data.best_model_path,
         target=ckpt,
         step=0,
         overwrite=True,
@@ -147,9 +144,8 @@ def test_ase_calc(get_tmp_path):
     )
     ckpt = {"model": {"params": params}, "epoch": 0}
 
-    best_dir = model_config.data.best_model_path
     checkpoints.save_checkpoint(
-        ckpt_dir=best_dir,
+        ckpt_dir=model_config.data.best_model_path,
         target=ckpt,
         step=0,
         overwrite=True,