Update modeling_hf_nomic_bert.py

modeling_hf_nomic_bert.py

@@ -16,7 +16,7 @@ from einops import rearrange, repeat
 from transformers import GPT2Config, PreTrainedModel
 from transformers.models.bert.modeling_bert import (
     BaseModelOutputWithPoolingAndCrossAttentions,
-    BertForPreTrainingOutput,
+    MaskedLMOutput,
     SequenceClassifierOutput
 )
 
@@ -1147,9 +1147,11 @@ class NomicBertForPreTraining(NomicBertPreTrainedModel):
             )
             total_loss = masked_lm_loss.float()
 
-        return BertForPreTrainingOutput(
+        return MaskedLMOutput(
             loss=total_loss,
-            prediction_logits=prediction_scores,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=None, 
         )