fix test

Signed-off-by: jiang1.li <[email protected]>
vllm-project · Nov 20, 2024 · 504c82b · 504c82b
1 parent 3fe3724
commit 504c82b
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/tests/basic_correctness/test_chunked_prefill.py b/tests/basic_correctness/test_chunked_prefill.py
@@ -282,7 +282,6 @@ def test_models_cpu(
     max_tokens: int,
     chunked_prefill_token_size: int,
     enforce_eager: bool,
-    tensor_parallel_size: int,
     attention_backend: str,
     monkeypatch,
 ) -> None:
@@ -295,7 +294,7 @@ def test_models_cpu(
         max_tokens,
         chunked_prefill_token_size,
         enforce_eager,
-        tensor_parallel_size,
+        1,
         attention_backend,
         monkeypatch,
     )
@@ -312,14 +311,13 @@ def test_with_prefix_caching_cpu(
     max_tokens: int,
     enforce_eager: bool,
     chunk_size: int,
-    tensor_parallel_size: int,
     dtype: str,
 ) -> None:
     test_with_prefix_caching(
         vllm_runner,
         max_tokens,
         enforce_eager,
         chunk_size,
-        tensor_parallel_size,
+        1,
         dtype,
     )