Update modeling_custom.py

modeling_custom.py

@@ -147,7 +147,7 @@ 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:
+        if input_ids.shape[0] == 1 and len(input_ids.shape) == 2 and input_ids[0,0] == input_ids[0,1] == 2:
              input_ids = input_ids[:, 1:]
              if attention_mask is not None:
                  attention_mask = attention_mask[:, 1:]