From 4726073865968980e40f22197d4747023455726d Mon Sep 17 00:00:00 2001
From: Ke Bao <ispobaoke@gmail.com>
Date: Mon, 2 Mar 2026 14:55:15 +0800
Subject: [PATCH] Fix mamba2 mixer ci test (#19658)

---
 test/registered/layers/mamba/test_mamba2_mixer.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test/registered/layers/mamba/test_mamba2_mixer.py b/test/registered/layers/mamba/test_mamba2_mixer.py
index 1bac53aca..bbc145103 100644
--- a/test/registered/layers/mamba/test_mamba2_mixer.py
+++ b/test/registered/layers/mamba/test_mamba2_mixer.py
@@ -40,7 +40,9 @@ def test_mixer2_gated_norm_multi_gpu(
     if not torch.cuda.is_available():
         pytest.skip("CUDA device not available")
 
-    assert torch.cuda.device_count() == NUM_GPUS
+    assert (
+        torch.cuda.device_count() >= NUM_GPUS
+    ), f"This test requires at least {NUM_GPUS} GPUs, but only {torch.cuda.device_count()} available"
 
     hidden_size, n_groups = hidden_size_n_groups
     num_processes = NUM_GPUS