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

This commit is contained in:
Woosuk Kwon 2024-06-28 18:14:16 -07:00 committed by GitHub
parent 7041de4384
commit 54814fd85b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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,