optimize dlinfer moe (#2741)

InternLM · Nov 13, 2024 · 7250318 · 7250318
1 parent 20544d3
commit 7250318
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/lmdeploy/pytorch/backends/dlinfer/moe.py b/lmdeploy/pytorch/backends/dlinfer/moe.py
@@ -19,9 +19,8 @@ def __init__(self, top_k: int, dim: int = -1):
 
     def forward(self, x: torch.Tensor):
         routing_weights, selected_experts = moe_gating_topk_softmax(
-            x, self.top_k)
-        return routing_weights.to(torch.float32), selected_experts.to(
-            torch.int64)
+            x.to(torch.float32), self.top_k)
+        return routing_weights, selected_experts
 
 
 class DlinferSoftmaxTopKBuilder(SoftmaxTopKBuilder):