adds pooling param

Signed-off-by: Flavia Beo <[email protected]>
flaviabeo · Oct 21, 2024 · d371268 · d371268
1 parent 180aad0
commit d371268
Showing 1 changed file with 3 additions and 4 deletions.
diff --git 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)