Add higher-capacity last-four-layer fusion

2026-08-26 13:06:18 +08:00
parent 43e0e14cf5
commit f45ea9cbc6
9 changed files with 397 additions and 17 deletions
@@ -979,6 +979,9 @@ def _save_weight_checkpoint(
                "num_horizons": NUM_HORIZONS,
                "num_classes": NUM_CLASSES,
                "dropout": raw_head.dropout,
+                "num_input_layers": raw_head.num_input_layers,
+                "layer_fusion": raw_head.layer_fusion,
+                "learnable_layer_mixture": raw_head.learnable_layer_mixture,
            },
            "initial_head_checkpoint": (
                str(args.head_checkpoint.resolve()) if args.head_checkpoint else None