moe test script

pytorch/cuda/moe_test.py

@@ -32,7 +32,9 @@ def perf():
     maxt = 0.
     sqtot = 0.
     for i in range(n_runs):
-        gate = torch.randint(low=0, high=num_expert, size=(batch_size, ), requires_grad=False).int().cuda()
+        gate = torch.randint(low=0, 
+                high=num_expert * torch.distributed.get_world_size(), 
+                size=(batch_size, ), requires_grad=False).int().cuda()
         ts = time.time()
         o = moe(inp, gate)
         te = time.time()