update unit test

33fc3aca · Rick Ho · d9dba929 · 33fc3aca
Commit 33fc3aca authored Oct 10, 2021 by Rick Ho
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

tests/test_numerical.py tests/test_numerical.py +2 -1

No files found.
--- a/tests/test_numerical.py
+++ b/tests/test_numerical.py
@@ -342,7 +342,8 @@ def _test_fmoe_local_ddp(rank, world_size, mp_group, dp_group, world_group):
    torch.cuda.manual_seed(42 + rank)

    model = MyModule().cuda()
-    model_ddp = LocalDDP(deepcopy(model), mp_group, dp_group, world_group)
+    model_ddp = LocalDDP(deepcopy(model),
+            mp_group=mp_group, dp_group=dp_group, world_group=world_group)
    model.set_comm()
    model_ddp.module.set_comm()