From e2c7dedb3a470b36463138cf0546f330b81dedad Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Thu, 25 Apr 2024 03:28:53 +0000
Subject: [PATCH] Minor

---
 vllm/worker/tpu_model_runner.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/vllm/worker/tpu_model_runner.py b/vllm/worker/tpu_model_runner.py
index f10d512791a54..fc2aab0261576 100644
--- a/vllm/worker/tpu_model_runner.py
+++ b/vllm/worker/tpu_model_runner.py
@@ -193,7 +193,7 @@ class TPUModelRunner:
             kv_caches,
             logits_indices,
         )
-        # TODO
+        # TODO(woosuk): Support sampling with temperature and top_p.
         next_token_ids = jnp.argmax(logits, axis=-1)
         return next_token_ids, new_kv_caches