From c2382b516cf8055c4c0d580fef9296dd5a409d0a Mon Sep 17 00:00:00 2001
From: Wooyeon Lee <wooyeon0.lee@samsung.com>
Date: Mon, 15 Jul 2024 18:47:45 +0900
Subject: [PATCH] mark. need 4 gpus

---
 tests/spec_decode/e2e/test_integration_dist_tp4.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/spec_decode/e2e/test_integration_dist_tp4.py b/tests/spec_decode/e2e/test_integration_dist_tp4.py
index a720e75717ea3..49e4a5f8150b5 100644
--- 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",
     [{