Add learnable last-four-layer query fusion

2026-08-26 10:29:43 +08:00
parent dbef1fc173
commit 43e0e14cf5
7 changed files with 268 additions and 37 deletions
@@ -726,12 +726,7 @@ def _load_initial_head(
        raise ValueError(
            f"Head hidden size {config['hidden_size']} != cache hidden size {hidden_size}"
        )
-    head = TurnPredictionHead(
-        hidden_size=hidden_size,
-        num_horizons=int(config["num_horizons"]),
-        num_classes=int(config["num_classes"]),
-        dropout=float(config["dropout"]),
-    )
+    head = TurnPredictionHead(**config)
    head.load_state_dict(payload.get("head", payload))
    if verbose:
        print(f"Initialized classification head from {checkpoint}")