fix

pytorch/cuda/moe.cpp

@@ -53,7 +53,7 @@ std::vector<torch::Tensor> moe_backward(
         Wx+b = [W b] [x]
                      [1]  
     */
-    return moe_cuda_forward(input, gate, weight);
+    return moe_cuda_backward(grad_output, input, gate, weight);
 }
 
 

pytorch/cuda/moe.py

@@ -49,3 +49,7 @@ input = torch.rand(batch_size, in_feat).cuda()
 gate = torch.randint(low=0, high=num_expert, size=(batch_size, )).int().cuda()
 
 output = moe(input, gate)
+
+
+y = output.mean()
+y.backward()
\ No newline at end of file