[CI/Build] Fix Args for _get_logits_warper in Sampler Test (vllm-pr…

…oject#5922)
opendatahub-io · Jul 1, 2024 · 98472e6 · 98472e6
1 parent c5ff677
commit 98472e6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/samplers/test_sampler.py b/tests/samplers/test_sampler.py
@@ -587,7 +587,7 @@ def test_sampler_top_k_top_p(seed: int, device: str):
     generation_config = GenerationConfig(top_k=top_k,
                                          top_p=top_p,
                                          do_sample=True)
-    warpers = generation_model._get_logits_warper(generation_config)
+    warpers = generation_model._get_logits_warper(generation_config, device)
     assert len(warpers) == 2  # top_p and top_k
 
     seq_group_metadata_list: List[SequenceGroupMetadata] = []