update flax and ipsuite tests

apax-hub · Dec 14, 2024 · a5e382b · a5e382b
1 parent 7071728
commit a5e382b
Show file tree

Hide file tree

Showing 4 changed files with 499 additions and 504 deletions.
diff --git a/apax/train/checkpoints.py b/apax/train/checkpoints.py
@@ -69,12 +69,13 @@ def create_params(model, rng_key, sample_input: tuple, n_models: int):
 
 
 def load_state(state, ckpt_dir):
+    ckpt_dir = Path(ckpt_dir)
     start_epoch = 0
     target = {"model": state, "epoch": 0}
-    checkpoints_exist = Path(ckpt_dir).is_dir()
+    checkpoints_exist = ckpt_dir.is_dir()
     if checkpoints_exist:
         log.info("Loading checkpoint")
-        raw_restored = checkpoints.restore_checkpoint(ckpt_dir, target=target, step=None)
+        raw_restored = checkpoints.restore_checkpoint(ckpt_dir.resolve(), target=target, step=None)
         state = raw_restored["model"]
         start_epoch = raw_restored["epoch"] + 1
         log.info("Successfully restored checkpoint from epoch %d", raw_restored["epoch"])