bugfix save and load ckpt from model_tag dir

2025-12-21 15:07:04 +08:00
parent d5759400f9
commit 92c6654b95
3 changed files with 14 additions and 13 deletions
@@ -207,7 +207,7 @@ while True:

    # save checkpoint at the end of the run (only on master process)
    if master_process and last_step and not dry_run:
-        output_dirname = f"d{depth}" # e.g. d12
+        output_dirname = model_tag if model_tag else f"d{depth}" # e.g. d12
        checkpoint_dir = os.path.join(base_dir, "mid_checkpoints", output_dirname)
        save_checkpoint(
            checkpoint_dir,