[Bugfix][TPU] Fix TPU sampler output (#5978)

2024-06-28 18:14:16 -07:00
parent 7041de4384
commit 54814fd85b
1 changed files with 1 additions and 1 deletions
--- a/vllm/worker/tpu_worker.py
+++ b/vllm/worker/tpu_worker.py
@ -215,7 +215,7 @@ class TPUWorker(LoraNotSupportedWorkerBase):
        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,