eval: strip torch.compile _orig_mod./module. prefixes when loading checkpoints (salvages compiled-GPU runs)

1 files changed, 11 insertions, 1 deletions

diff --git a/src/eval.py b/src/eval.py
index 191ac79..5c329a6 100644
--- a/src/eval.py
+++ b/src/eval.py
@@ -243,7 +243,17 @@ def _load(out_dir: str, ckpt: str, cfg: Config, device: torch.device | None = No
         device_str = getattr(cfg, "device", "auto")
         device = torch.device("cuda" if (device_str == "cuda" or (device_str == "auto" and torch.cuda.is_available())) else "cpu")
     model = build_model(cfg).to(device)
-    model.load_state_dict(torch.load(os.path.join(out_dir, ckpt), map_location=device))
+    state = torch.load(os.path.join(out_dir, ckpt), map_location=device)
+    # torch.compile wraps modules (model._orig_mod), so checkpoints saved from a
+    # compiled model carry an "_orig_mod." key prefix; DataParallel adds "module.".
+    # Strip both so the uncompiled eval build loads cleanly.
+    clean = {}
+    for k, v in state.items():
+        for pfx in ("_orig_mod.", "module."):
+            while k.startswith(pfx):
+                k = k[len(pfx):]
+        clean[k] = v
+    model.load_state_dict(clean)
     model.eval()
     return model