Merge pull request #200 from Heihaierr/master

Update switch_gate.py

Author: laekov

fmoe/gates/switch_gate.py

@@ -42,7 +42,7 @@ def forward(self, inp):
         top1_score = top1_score.to(dtype=inp.dtype)
 
         cap_rate = self.capacity[0 if self.training else 1]
-        capacity = math.ceil(cap_rate * inp.shape[0])
+        capacity = math.ceil(cap_rate * inp.shape[0] / self.num_expert)
         _new_lec, _new_gec, top1_idx = limit_by_capacity(
                 top1_idx, self.num_expert, self.world_size, capacity)