From 7be3a59d8ee7014d6462c258222cbfa8be815831 Mon Sep 17 00:00:00 2001 From: Ning Xie Date: Mon, 18 Aug 2025 13:09:08 +0800 Subject: [PATCH] [Misc] enhance static type hint (#23059) Signed-off-by: Andy Xie --- vllm/v1/worker/lora_model_runner_mixin.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/lora_model_runner_mixin.py b/vllm/v1/worker/lora_model_runner_mixin.py index 2fbdee4724e3..84ed46989ea9 100644 --- a/vllm/v1/worker/lora_model_runner_mixin.py +++ b/vllm/v1/worker/lora_model_runner_mixin.py @@ -8,6 +8,7 @@ from contextlib import contextmanager from typing import Union import numpy as np +import torch import torch.nn as nn from vllm.config import LoRAConfig, ModelConfig, SchedulerConfig @@ -31,7 +32,8 @@ class LoRAModelRunnerMixin: def load_lora_model(self, model: nn.Module, model_config: ModelConfig, scheduler_config: SchedulerConfig, - lora_config: LoRAConfig, device: str) -> nn.Module: + lora_config: LoRAConfig, + device: torch.device) -> nn.Module: if not supports_lora(model): raise ValueError(