Update modeling_esm_plusplus.py

modeling_esm_plusplus.py

@@ -669,7 +669,12 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM):
         Returns:
             ESMplusplusOutput containing loss, logits, and hidden states
         """
-        output = super().forward(input_ids, attention_mask, labels, output_hidden_states)
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=None,
+            output_hidden_states=output_hidden_states
+        )
         x = output.last_hidden_state
         cls_features = x[:, 0, :]
         mean_features = self.mean_pooling(x, attention_mask)
@@ -735,7 +740,12 @@ class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM):
         Returns:
             ESMplusplusOutput containing loss, logits, and hidden states
         """
-        output = super().forward(input_ids, attention_mask, labels, output_hidden_states)
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=None,
+            output_hidden_states=output_hidden_states
+        )
         x = output.last_hidden_state
         logits = self.classifier(x)
         loss = None