fix zero test

07a5d8ac · Rick Ho · 5f8ba136 · 07a5d8ac
Commit 07a5d8ac authored Oct 10, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 2 deletions

tests/test_zero.py tests/test_zero.py +3 -2

No files found.
--- a/tests/test_zero.py
+++ b/tests/test_zero.py
@@ -56,8 +56,9 @@ def _test_zero_transformer(num_expert=2, batch_size=4, d_hidden=8, world_size=1)
    mask_dict = {
        1: torch.zeros(d_hidden).cuda()
    }
-    model = FMoETransformerMLP(num_expert, d_hidden, d_hidden * 4, world_size,
-            gate=ConstantGate, mask=mask, mask_dict=mask_dict).cuda()
+    model = FMoETransformerMLP(num_expert, d_hidden, d_hidden * 4,
+            world_size=world_size, gate=ConstantGate, mask=mask,
+            mask_dict=mask_dict).cuda()
    oup = model(inp)