Add return_dict to NomicBertModel; required for Sentence Transformers

modeling_hf_nomic_bert.py

@@ -1069,6 +1069,7 @@ class NomicBertModel(NomicBertPreTrainedModel):
         position_ids=None,
         token_type_ids=None,
         attention_mask=None,
+        return_dict=None,
     ):
         if token_type_ids is None:
             token_type_ids = torch.zeros_like(input_ids)
@@ -1080,7 +1081,7 @@ class NomicBertModel(NomicBertPreTrainedModel):
 
         attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape)
         sequence_output = self.encoder(
-            hidden_states, attention_mask=attention_mask
+            hidden_states, attention_mask=attention_mask, return_dict=return_dict,
         )
 
         pooled_output = self.pooler(sequence_output) if self.pooler is not None else None