From 91b3929605cdfb6cd8536ece705d82843418e2e0 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Mon, 1 Jun 2026 01:14:38 +0000 Subject: [PATCH] fix: call moe_runner.run() and se_runner.run() (not __call__) --- single_shot_inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/single_shot_inference.py b/single_shot_inference.py index 614dfc32..de1ea9a6 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -380,7 +380,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, # ===================================================================== def moe_forward(x, li, moe_runner, se_runner, router, token_id): topk_w, topk_ids = router(x, token_ids=token_id) - routed_out = moe_runner(x, topk_w, topk_ids); shared_out = se_runner(x) + routed_out = moe_runner.run(x, topk_w, topk_ids) + shared_out = se_runner.run(x) return routed_out + shared_out # =====================================================================