diff --git a/vllm/model_executor/models/bailing_moe_linear.py b/vllm/model_executor/models/bailing_moe_linear.py index 36e91129bae..c49e646c389 100644 --- a/vllm/model_executor/models/bailing_moe_linear.py +++ b/vllm/model_executor/models/bailing_moe_linear.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -from collections.abc import Iterable -from typing import Any +from collections.abc import Callable, Iterable import torch import torch.nn as nn @@ -618,7 +617,9 @@ class BailingMoeV25Model(nn.Module): return False param = params_dict[name] - weight_loader: Any = getattr(param, "weight_loader", default_weight_loader) + weight_loader: Callable[..., None] = getattr( + param, "weight_loader", default_weight_loader + ) if shard_id is None: weight_loader(param, tensor) diff --git a/vllm/model_executor/models/bailing_moe_mtp.py b/vllm/model_executor/models/bailing_moe_mtp.py index a827faef633..a403244a177 100644 --- a/vllm/model_executor/models/bailing_moe_mtp.py +++ b/vllm/model_executor/models/bailing_moe_mtp.py @@ -2,8 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Inference-only Bailing MoE v2.5 MTP model.""" -from collections.abc import Iterable -from typing import Any +from collections.abc import Callable, Iterable import torch import torch.nn as nn @@ -277,7 +276,9 @@ class BailingMoeV25MTPModel(nn.Module): return False param = params_dict[name] - weight_loader: Any = getattr(param, "weight_loader", default_weight_loader) + weight_loader: Callable[..., None] = getattr( + param, "weight_loader", default_weight_loader + ) if shard_id is None: weight_loader(param, loaded_weight) elif isinstance(shard_id, int):