Fix mamba2 mixer ci test (#19658)

2026-03-02 14:55:15 +08:00
parent 922aad2faa
commit 4726073865
1 changed files with 3 additions and 1 deletions
@@ -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