From f7765ede6f738a6223ac154911e19e699edaa9b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fl=C3=A1via=20B=C3=A9o?= <119421251+flaviabeo@users.noreply.github.com> Date: Wed, 23 Oct 2024 13:35:14 -0300 Subject: [PATCH] Update vllm/model_executor/layers/pooler.py default Co-authored-by: Cyrus Leung --- vllm/model_executor/layers/pooler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/pooler.py b/vllm/model_executor/layers/pooler.py index d7cb111742836..221bb77434868 100644 --- a/vllm/model_executor/layers/pooler.py +++ b/vllm/model_executor/layers/pooler.py @@ -40,7 +40,7 @@ def get_pooling_type(self, pooling_type_name: str) -> PoolingType: pooling_types = PoolingType.__dict__.items() return PoolingType( next((value for key, value in pooling_types - if key.lower() in pooling_type_name), 2)) + if key.lower() in pooling_type_name), PoolingType.CLS)) class Pooler(nn.Module):