fix: huggingface#14486 do not use BertPooler in DPR

PaulLerner · Jan 7, 2022 · 95eaf44 · 95eaf44
1 parent f18c6fa
commit 95eaf44
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/models/dpr/modeling_dpr.py b/src/transformers/models/dpr/modeling_dpr.py
@@ -175,7 +175,7 @@ class DPREncoder(DPRPreTrainedModel):
 
     def __init__(self, config: DPRConfig):
         super().__init__(config)
-        self.bert_model = BertModel(config)
+        self.bert_model = BertModel(config, add_pooling_layer=False)
         assert self.bert_model.config.hidden_size > 0, "Encoder hidden_size can't be zero"
         self.projection_dim = config.projection_dim
         if self.projection_dim > 0:
@@ -202,7 +202,7 @@ def forward(
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
         )
-        sequence_output, pooled_output = outputs[:2]
+        sequence_output = outputs[0]
         pooled_output = sequence_output[:, 0, :]
         if self.projection_dim > 0:
             pooled_output = self.encode_proj(pooled_output)