From 1916e0f5998705ee89f5971b4a8b9fb1a040e21f 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,