Update custom_model_package/custom_model.py

custom_model_package/custom_model.py

@@ -42,12 +42,13 @@ class CustomModel(XLMRobertaForSequenceClassification):
         with torch.no_grad():
             cls_token_state = sequence_output[:, 0, :].unsqueeze(1)
             sentiment_logits = self.classifier(cls_token_state).squeeze(1)
+        logits = torch.cat([sentiment_logits, emotion_logits], dim=-1)
         if labels is not None:
             class_weights = torch.tensor([1.0] * self.num_emotion_labels).to(labels.device)
             loss_fct = nn.BCEWithLogitsLoss(pos_weight=class_weights)
             loss = loss_fct(emotion_logits, labels)
-            return {"loss": loss, "emotion_logits": emotion_logits, "sentiment_logits": sentiment_logits}
-        return {"emotion_logits": emotion_logits, "sentiment_logits": sentiment_logits}
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):