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/evaluate_single_disease_mortality_attribution.py
+++ b/evaluate_single_disease_mortality_attribution.py
@@ -553,7 +553,9 @@ def main() -> None:
    device = resolve_eval_device(args.device)
    selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
    selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
-    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()