Return BaseModelOutput

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 .bert_padding import index_put_first_axis
 
 from .activation import get_act_fn
@@ -874,7 +874,7 @@ class FlexBertPreTrainedModel(BertPreTrainedModel):
             init_weights(self.config, module, type_of_module=ModuleType.emb)
         else:
             print("Custom weight init for the given module is not supported")
-            print(module)
+            # print(module)
             # raise NotImplementedError("Custom weight init for the given module is not supported")
 
 
@@ -967,7 +967,8 @@ class FlexBertModel(FlexBertPreTrainedModel):
 
         if self.final_norm is not None:
             encoder_outputs = self.final_norm(encoder_outputs)
-        return encoder_outputs
+        
+        return BaseModelOutput(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"