Merge pull request #43 from haohongxiang/fix_bug_of_recompute

fix bug of recompute in moe_layer
PaddlePaddle · Feb 24, 2022 · 3f27dea · 3f27dea
2 parents 91ba8b3 + 9ff23f5
commit 3f27dea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/paddle/distributed/model/moe/moe_layer.py b/python/paddle/distributed/model/moe/moe_layer.py
@@ -393,7 +393,7 @@ def experts_fwd(x, fwd_expert_count, experts):
             else:
                 return paddle.concat(y, axis=0)
 
-        if self.recompute_interval <= 0:
+        if self.recompute_interval <= 0 or x.shape[0] == 0:
             x = experts_fwd(x, fwd_expert_count.numpy(), self.experts)
         else:
             x = _hp_recompute(experts_fwd, x,