[TPU] Fix the test_sampler (#17820)

2025-05-08 02:51:33 -07:00
parent ca04b97c93
commit a463555dee
2 changed files with 2 additions and 2 deletions
--- a/tests/v1/tpu/test_sampler.py
+++ b/tests/v1/tpu/test_sampler.py
@@ -26,7 +26,7 @@ def test_sampler_different(model_name: str):
              enforce_eager=False,
              max_num_seqs=1,
              max_model_len=512,
-              max_num_batched_tokens=512)
+              max_num_batched_tokens=256)
    prompts = [
        "Write a short story about a robot that dreams for the first time."
    ]