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

5d5b4c5f · Woosuk Kwon · GitHub · 70c232f8 · 5d5b4c5f
Unverified Commit 5d5b4c5f authored Jul 09, 2024 by Woosuk Kwon Committed by GitHub Jul 09, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

vllm/worker/tpu_model_runner.py vllm/worker/tpu_model_runner.py +1 -1

No files found.
--- 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,