feat: added return_dict

modeling_bert.py

@@ -379,6 +379,7 @@ class BertModel(BertPreTrainedModel):
         task_type_ids=None,
         attention_mask=None,
         masked_tokens_mask=None,
+        return_dict=True,
     ):
         """If masked_tokens_mask is not None (i.e. last_layer_subset == True in BertForPreTraining),
         we only want the output for the masked tokens. This means that we only compute the last
@@ -429,6 +430,9 @@ class BertModel(BertPreTrainedModel):
                 sequence_output = sequence_output[masked_tokens_mask[subset_mask]]
             pooled_output = self.pooler(pool_input, pool=False) if self.pooler is not None else None
 
+        if not return_dict:
+            return (sequence_output, pooled_output)
+        
         return BaseModelOutputWithPoolingAndCrossAttentions(
             last_hidden_state=sequence_output,
             pooler_output=pooled_output,