Unify FFN and TrajMixer model architectures

2026-07-25 12:59:09 +08:00
parent 4526191fe1
commit b13db5e407
18 changed files with 1019 additions and 52 deletions
--- a/export_weibull_death_parameter_stats.py
+++ b/export_weibull_death_parameter_stats.py
@@ -381,7 +381,9 @@ def main() -> None:
    cfg_model = dict(cfg)
    cfg_model["dist_mode"] = dist_mode
    device = resolve_eval_device(args.device)
-    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
    load_model_state(model, state_dict)
    model.eval()