[Bugfix][TPU] Fix TPU sampler output (vllm-project#5978)

opendatahub-io · Jul 1, 2024 · 28773bf · 28773bf
1 parent 96e23ec
commit 28773bf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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,