From d371268f537a017c66ae3dffb2317247667a13dc Mon Sep 17 00:00:00 2001 From: Flavia Beo Date: Mon, 21 Oct 2024 11:07:39 -0300 Subject: [PATCH] adds pooling param Signed-off-by: Flavia Beo --- vllm/model_executor/model_loader/loader.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/model_loader/loader.py b/vllm/model_executor/model_loader/loader.py index 31d0210f5aaa7..f0ecdc61a3fed 100644 --- a/vllm/model_executor/model_loader/loader.py +++ b/vllm/model_executor/model_loader/loader.py @@ -145,8 +145,6 @@ def _get_model_initialization_kwargs( if has_inner_state(model_class) and scheduler_config: extra_kwargs["scheduler_config"] = scheduler_config - extra_kwargs["pooling_type"] = pooling_type - return extra_kwargs @@ -159,12 +157,13 @@ def build_model(model_class: Type[nn.Module], hf_config: PretrainedConfig, pooling_type) -> nn.Module: extra_kwargs = _get_model_initialization_kwargs(model_class, lora_config, multimodal_config, - scheduler_config, - pooling_type) + scheduler_config + ) return model_class(config=hf_config, cache_config=cache_config, quant_config=quant_config, + pooling_type **extra_kwargs)