Return BaseModelOutputWithPoolingAndCrossAttentions

modeling_flexbert.py

@@ -66,7 +66,7 @@ from transformers.modeling_outputs import (
     SequenceClassifierOutput,
 )
 from transformers.models.bert.modeling_bert import BertPreTrainedModel
-from transformers import BaseModelOutput
+from transformers import BaseModelOutputWithPoolingAndCrossAttentions
 from .bert_padding import index_put_first_axis
 
 from .activation import get_act_fn
@@ -968,7 +968,7 @@ class FlexBertModel(FlexBertPreTrainedModel):
         if self.final_norm is not None:
             encoder_outputs = self.final_norm(encoder_outputs)
         
-        return BaseModelOutput(last_hidden_state=encoder_outputs)
+        return BaseModelOutputWithPoolingAndCrossAttentions(last_hidden_state=encoder_outputs)
 
     def _init_weights(self, module: Optional[nn.Module] = None, reset_params: Optional[bool] = None):
         assert (module is None) != (reset_params is None), "arg module xor reset_params must be specified"