update models

inference.py

@@ -238,8 +238,30 @@ def build_torch_model_from_ckpt(model_name: str, ckpt: dict, device: torch.devic
         model = CNNHead(in_ch=in_dim, c=int(params["channels"]), k=int(params["kernel"]),
                         layers=int(params["layers"]), dropout=dropout)
     elif model_name == "transformer":
-        model = TransformerHead(in_dim=in_dim, d_model=int(params["d_model"]), nhead=int(params["nhead"]),
-                                layers=int(params["layers"]), ff=int(params["ff"]), dropout=dropout)
+        print(
+            f"[LOAD] {prop_key}/{mode} transformer params keys:",
+            list(params.keys())
+        )
+
+        d_model = (
+            params.get("d_model")
+            or params.get("hidden")
+            or params.get("hidden_dim")
+        )
+        if d_model is None:
+            raise KeyError(
+                f"Transformer checkpoint missing d_model/hidden. "
+                f"Available keys: {list(params.keys())}"
+            )
+
+        model = TransformerHead(
+            in_dim=in_dim,
+            d_model=int(d_model),
+            nhead=int(params["nhead"]),
+            layers=int(params["layers"]),
+            ff=int(params.get("ff", 4 * int(d_model))),
+            dropout=dropout
+        )
     else:
         raise ValueError(f"Unknown NN model_name={model_name}")