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

54814fd8 · Woosuk Kwon · GitHub · 7041de43 · 54814fd8
Unverified Commit 54814fd8 authored Jun 28, 2024 by Woosuk Kwon Committed by GitHub Jun 28, 2024
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

vllm/worker/tpu_worker.py vllm/worker/tpu_worker.py +1 -1

No files found.
--- 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,