From fda9f75f350f751309a920466a1cda401603b129 Mon Sep 17 00:00:00 2001 From: Hao Zhang Date: Fri, 27 Mar 2026 10:34:59 +0800 Subject: [PATCH] chore: Fix a typo in the previous commit. --- qmp/networks/transformers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qmp/networks/transformers.py b/qmp/networks/transformers.py index d99537d..162949e 100644 --- a/qmp/networks/transformers.py +++ b/qmp/networks/transformers.py @@ -143,7 +143,7 @@ def forward( similarity = torch.nn.functional.softmax(x @ self.centroid.t(), dim=-1) # top_k_indices: batch * site * selected _, top_k_indices = torch.topk(similarity + self.bias, self.selected_num, dim=-1) - # gate, gate: batch * site * routed + # gate: batch * site * routed gate = torch.zeros_like(similarity).scatter_(-1, top_k_indices, similarity.gather(-1, top_k_indices)) for i, expert in enumerate(self.feed_forward_routed): y = y + expert(x) * gate[:, :, i].unsqueeze(-1)