Skip to content

GitLab

Unverified Commit 63e2a641 authored by Zhuohan Li's avatar Zhuohan Li Committed by GitHub
Browse files

[FIX] Fix beam search test (#2930)

parent 264017a2
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
tests/samplers/test_beam_search.py
View file @ 63e2a641
...@@ -26,6 +26,7 @@ def test_beam_search_single_input( ...@@ -26,6 +26,7 @@ def test_beam_search_single_input(
max_tokens: int, max_tokens: int,
beam_width: int, beam_width: int,
) -> None: ) -> None:
example_prompts = example_prompts[:1]
hf_model = hf_runner(model, dtype=dtype) hf_model = hf_runner(model, dtype=dtype)
hf_outputs = hf_model.generate_beam_search(example_prompts, beam_width, hf_outputs = hf_model.generate_beam_search(example_prompts, beam_width,
max_tokens) max_tokens)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment