Update vllm/model_executor/layers/pooler.py default

Co-authored-by: Cyrus Leung <[email protected]>
flaviabeo · Oct 23, 2024 · f7765ed · f7765ed
1 parent 349b3f2
commit f7765ed
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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):