Fix loss function

modeling_hf_alibaba_nlp_gte.py

@@ -975,7 +975,6 @@ class GteForSequenceClassification(GtePreTrainedModel):
         self.model = GteModel(config, add_pooling_layer=config.add_pooling_layer)
 
         self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
-        self.loss_function = nn.MSELoss()
         self.post_init()
 
     def get_input_embeddings(self):
@@ -1020,7 +1019,7 @@ class GteForSequenceClassification(GtePreTrainedModel):
 
         loss = None
         if labels is not None:
-            loss = self.loss_function(logits.squeeze(-1), labels.squeeze(-1))
+            loss = self.loss_function(labels, logits, self.config)
 
         return SequenceClassifierOutputWithPast(
             loss=loss,