fix args

vllm-project · Aug 28, 2024 · a0bac86 · a0bac86
1 parent 80ce2bd
commit a0bac86
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/vllm/worker/tpu_model_runner.py b/vllm/worker/tpu_model_runner.py
@@ -232,8 +232,15 @@ def _dummy_run(
             torch._dynamo.mark_dynamic(t, 0)
             torch._dynamo.mark_dynamic(p, 0)
         # Dummy run.
-        self.model(token_ids, position_ids, attn_metadata, input_lens, t, p,
-                   num_samples, kv_caches)
+        self.model(token_ids,
+                   position_ids,
+                   attn_metadata,
+                   input_lens,
+                   t,
+                   p,
+                   num_samples,
+                   kv_caches,
+                   is_prompt=is_prompt)
 
     def warmup_model(
         self,