Update modeling_custom.py

modeling_custom.py

@@ -148,6 +148,11 @@ class Gemma2ForQuantileSequenceClassification(Gemma2PreTrainedModel):
             `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
         """
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if input_ids.shape[0] == 1 and len(input_ids.shape) == 2 and input_ids[0,0] == input_ids[0,0] == 2:
+             input_ids = input_ids[:, 1:]
+             if attention_mask is not None:
+                 attention_mask = attention_mask[:, 1:]
+                 
         transformer_outputs = self.model(
             input_ids,
             attention_mask=attention_mask,