update fused_moe.py

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -740,7 +740,7 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
             expert_ids,
             num_tokens_post_padded,
             B.shape[1] if not use_nn_moe else B.shape[2],
-            A.shape[1] if not use_nn_moe else A.shape[2],
+            A.shape[1],
             EM,
             topk_ids.numel(),
             A.stride(0),