Upload modelling_xlm_roberta.py

modelling_xlm_roberta.py

@@ -1456,8 +1456,9 @@ class XLMRobertaForEstimation(XLMRobertaPreTrainedModel):
             return_dict=return_dict,
         )
 
-        logits = self.estimator(outputs[0])
-        reshaped_logits = logits.view(-1, num_choices)
+        CLS_tok = outputs[0][:, 0, :]  # for some reason at sentence level we take the first token score cf Comet
+        logits = self.estimator(CLS_tok)
+        reshaped_logits = logits #.view(-1, num_choices)
 
         loss = None
         if labels is not None: