feat: add scoliosis skeleton experiment tooling

2026-03-10 15:03:53 +08:00
parent 2647398307
commit 44e62ae3ae
6 changed files with 39 additions and 14 deletions
@@ -353,6 +353,18 @@ class BaseModel(MetaModel, nn.Module):
            weights_only=False,
        )
        model_state_dict = checkpoint['model']
+        current_state_keys = set(self.state_dict().keys())
+        stale_loss_keys = [
+            key for key in model_state_dict.keys()
+            if key.startswith("loss_aggregator.") and key not in current_state_keys
+        ]
+        for key in stale_loss_keys:
+            model_state_dict.pop(key)
+        if stale_loss_keys:
+            self.msg_mgr.log_warning(
+                "Ignoring stale loss state from %s: %s"
+                % (save_name, stale_loss_keys)
+            )

        if not load_ckpt_strict:
            self.msg_mgr.log_info("-------- Restored Params List --------")