diff --git a/tests/samplers/test_beam_search.py b/tests/samplers/test_beam_search.py index aef7eec09..98675856a 100644 --- a/tests/samplers/test_beam_search.py +++ b/tests/samplers/test_beam_search.py @@ -27,7 +27,7 @@ EXTRA_ENGINE_KWARGS: dict = ( max_num_seqs=1, ) if current_platform.is_rocm() - else {} + else dict(async_scheduling=False, max_num_seqs=1) ) # FIXME(zhuohan): The test can not pass if we: