return dict

modeling_hf_nomic_bert.py

@@ -1058,6 +1058,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)
@@ -1066,7 +1067,7 @@ class NomicBertModel(NomicBertPreTrainedModel):
         hidden_states = self.emb_drop(hidden_states)
 
         attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape)
-        sequence_output = self.encoder(hidden_states, attention_mask=attention_mask)
+        sequence_output = self.encoder(hidden_states, attention_mask=attention_mask, return_dict=return_dict)
 
         pooled_output = self.pooler(sequence_output) if self.pooler is not None else None