Unify FFN and TrajMixer model architectures

This commit is contained in:
2026-07-25 12:59:09 +08:00
parent 4526191fe1
commit b13db5e407
18 changed files with 1019 additions and 52 deletions

View File

@@ -649,7 +649,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()