[CI/Build] Fix Args for `_get_logits_warper` in Sampler Test (#5922)

736ed388 · Roger Wang · GitHub · 365791ff · 736ed388
Unverified Commit 736ed388 authored Jun 27, 2024 by Roger Wang Committed by GitHub Jun 27, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

tests/samplers/test_sampler.py tests/samplers/test_sampler.py +1 -1

No files found.
--- 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] = []