diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 14cfd00c2b..89e28d950f 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -1632,6 +1632,7 @@ def _parallel_worker( if all2all_manager is not None: all2all_manager.destroy() total = total + 1 + torch.distributed.barrier() skipped = total - (passed + failed)