From 9d69c6fb3e1ac6c858abeec3e6020c832ffbc924 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Mon, 22 Sep 2025 16:10:41 -0400 Subject: [PATCH] Fix --- scripts/benchmarks/deepseek_v3_moe.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/scripts/benchmarks/deepseek_v3_moe.py b/scripts/benchmarks/deepseek_v3_moe.py index 7385ca1e9..4d65b4e20 100644 --- a/scripts/benchmarks/deepseek_v3_moe.py +++ b/scripts/benchmarks/deepseek_v3_moe.py @@ -193,8 +193,9 @@ def main() -> None: # pragma: no cover - CLI entrypoint ) def _uniform_gate(self, hidden_states): - batch_tokens = hidden_states.shape[0] - return topk_idx[:batch_tokens], weights[:batch_tokens] + flat = hidden_states.view(-1, hidden_states.shape[-1]) + token_count = flat.shape[0] + return topk_idx[:token_count], weights[:token_count] patched_module.gate.forward = _uniform_gate.__get__( patched_module.gate, patched_module.gate.__class__