mark. need 4 gpus

vllm-project · Jul 15, 2024 · c2382b5 · c2382b5
1 parent 02dc475
commit c2382b5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/spec_decode/e2e/test_integration_dist_tp4.py b/tests/spec_decode/e2e/test_integration_dist_tp4.py
@@ -60,6 +60,8 @@ def test_draft_model_tp_lt_target_model_tp4(test_llm_generator,
                                          force_output_len=True)
 
 
+@pytest.mark.skipif(torch.cuda.device_count() < 4,
+                    reason="Need at least 4 GPUs to run the test.")
 @pytest.mark.parametrize(
     "common_llm_kwargs",
     [{