fix neuron

Signed-off-by: youkaichao <[email protected]>
vllm-project · Jan 14, 2025 · f5cf4e9 · f5cf4e9
1 parent 0be6c88
commit f5cf4e9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/vllm/worker/neuron_worker.py b/vllm/worker/neuron_worker.py
@@ -26,6 +26,7 @@ def __init__(
         local_rank: int,
         rank: int,
         distributed_init_method: str,
+        is_driver_worker: bool = True,
     ) -> None:
         WorkerBase.__init__(self, vllm_config=vllm_config)
         self.local_rank = local_rank
@@ -38,7 +39,7 @@ def __init__(
 
         self.model_runner: NeuronModelRunner = NeuronModelRunner(
             vllm_config=vllm_config)
-        self.is_driver_worker = True
+        self.is_driver_worker = is_driver_worker
 
     def execute_model(
         self,