[Bugfix] Fix test_mixtral_moe (#24371)

7555d6b3 · Jee Jee Li · GitHub · 00a4e56d · 7555d6b3
Unverified Commit 7555d6b3 authored Sep 07, 2025 by Jee Jee Li Committed by GitHub Sep 06, 2025
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

tests/kernels/moe/test_moe.py tests/kernels/moe/test_moe.py +2 -2

No files found.
--- a/tests/kernels/moe/test_moe.py
+++ b/tests/kernels/moe/test_moe.py
@@ -371,8 +371,8 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
 @pytest.mark.parametrize(
    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False])
 @torch.inference_mode()
-def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
+def test_mixtral_moe(dist_init, dtype: torch.dtype, padding: bool,
-                     monkeypatch):
+                     use_rocm_aiter: bool, monkeypatch):
    """Make sure our Mixtral MoE implementation agrees with the one from
    huggingface."""