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

This commit is contained in:
Woosuk Kwon 2024-07-09 00:21:37 -07:00 committed by GitHub
parent 70c232f85a
commit 5d5b4c5fe5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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,