[FIX] Fix beam search test (#2930)

2025-12-10 00:15:51 +08:00 · 2024-02-20 14:37:39 -08:00 · 2024-02-20 14:37:39 -08:00 · 63e2a6419d
commit 63e2a6419d
parent 264017a2bf
1 changed files with 1 additions and 0 deletions
--- a/tests/samplers/test_beam_search.py
+++ b/tests/samplers/test_beam_search.py
@ -26,6 +26,7 @@ def test_beam_search_single_input(
    max_tokens: int,
    beam_width: int,
 ) -> None:
+    example_prompts = example_prompts[:1]
    hf_model = hf_runner(model, dtype=dtype)
    hf_outputs = hf_model.generate_beam_search(example_prompts, beam_width,
                                               max_tokens)