[Bugfix][TPU] Add missing None to model input (#6245)

2024-07-09 00:21:37 -07:00
parent 70c232f85a
commit 5d5b4c5fe5
1 changed files with 1 additions and 1 deletions
--- a/vllm/worker/tpu_model_runner.py
+++ b/vllm/worker/tpu_model_runner.py
@ -154,7 +154,7 @@ class TPUModelRunner:
        # Dummy run.
        num_samples = _MAX_NUM_SAMPLES if is_prompt else 1
        self.model(token_ids, position_ids, kv_caches, attn_metadata,
-                   input_lens, t, p, num_samples)
+                   input_lens, None, t, p, num_samples)

    def warmup_model(
        self,