fix another bug to make global moe run

14c0eab4 · Rick Ho · 0fea2991 · 14c0eab4
Commit 14c0eab4 authored Jan 29, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

fmoe/fmoe_functions.py fmoe/fmoe_functions.py +1 -1

No files found.
--- a/fmoe/fmoe_functions.py
+++ b/fmoe/fmoe_functions.py
@@ -47,7 +47,7 @@ class MOEScatter(Function):
        (fwd_batch_size, local_batch_size, world_size) = ctx.moe_args

        if world_size > 1:
-            local_grad_in, = fmoe_cuda.global_gather(global_grad_out,
+            local_grad_in, = fmoe_cuda.global_gather(global_grad_in,
                    local_expert_count, global_expert_count,
                    local_batch_size, world_size)
        else: