Update modeling_esm_plusplus.py

modeling_esm_plusplus.py

@@ -411,6 +411,7 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM):
         if attention_mask is None:
             return x.mean(dim=1)
         else:
+            attention_mask = attention_mask.unsqueeze(-1)
             return (x * attention_mask).sum(dim=1) / attention_mask.sum(dim=1)
 
     def forward(
@@ -440,7 +441,7 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM):
 
             if self.config.problem_type == "regression":
                 if self.num_labels == 1:
-                    loss = self.mse(logits.squeeze(), labels.squeeze())
+                    loss = self.mse(logits.flatten(), labels.flatten())
                 else:
                     loss = self.mse(logits, labels)
             elif self.config.problem_type == "single_label_classification":