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