From 28773bf877e3be7e22daa537a83d2656a26c58aa Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 28 Jun 2024 18:14:16 -0700 Subject: [PATCH] [Bugfix][TPU] Fix TPU sampler output (#5978) --- vllm/worker/tpu_worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/worker/tpu_worker.py b/vllm/worker/tpu_worker.py index 37d810e8392a9..d58c7dc994184 100644 --- a/vllm/worker/tpu_worker.py +++ b/vllm/worker/tpu_worker.py @@ -215,7 +215,7 @@ def execute_model( assert len(seq_group_metadata_list) > 0 output = self.model_runner.execute_model(seq_group_metadata_list, self.tpu_cache) - return [output] + return output def cache_swap( self,