Fix

modeling_hf_alibaba_nlp_gte.py

@@ -16,7 +16,7 @@
 
 import math
 from dataclasses import dataclass
-from typing import List, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 import torch.utils.checkpoint
@@ -994,6 +994,7 @@ class GteForSequenceClassification(GtePreTrainedModel):
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
+        **kwargs: Any,
     ) -> SequenceClassifierOutputWithPast:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -1009,6 +1010,7 @@ class GteForSequenceClassification(GtePreTrainedModel):
             inputs_embeds=inputs_embeds,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
+            **kwargs
         )
         if self.config.add_pooling_layer:
             hidden_states = transformer_outputs.pooler_output
@@ -1021,6 +1023,10 @@ class GteForSequenceClassification(GtePreTrainedModel):
         if labels is not None:
             loss = self.loss_function(labels, logits, self.config)
 
+        # if not return_dict:
+        #     output = (logits,) + transformer_outputs[1:]
+        #     return ((loss,) + output) if loss is not None else output
+
         return SequenceClassifierOutputWithPast(
             loss=loss,
             logits=logits,