[BugFix] Fix return type of executor execute_model methods (#4402)

2024-04-27 11:17:45 -07:00
parent d6e520e170
commit ba4be44c32
6 changed files with 9 additions and 8 deletions
--- a/vllm/executor/ray_gpu_executor.py
+++ b/vllm/executor/ray_gpu_executor.py
@@ -188,7 +188,7 @@ class RayGPUExecutor(DistributedGPUExecutor):
                      blocks_to_swap_in: Dict[int, int],
                      blocks_to_swap_out: Dict[int, int],
                      blocks_to_copy: Dict[int, List[int]],
-                      num_lookahead_slots: int = 0) -> SamplerOutput:
+                      num_lookahead_slots: int = 0) -> List[SamplerOutput]:
        all_outputs = self._run_workers(
            "execute_model",
            driver_kwargs={