Update modeling_internlm2.py

modeling_internlm2.py

@@ -1968,7 +1968,7 @@ class InternLM2ForRewardModel(InternLM2PreTrainedModel):
         r_pad_attention_masks = torch.nn.utils.rnn.pad_sequence(attention_masks, batch_first=True, padding_value=False)
 
         outputs = self.forward(input_ids=r_pad_batch_input_ids.to(self.device), attention_mask=r_pad_attention_masks.to(self.device), **kwargs)
-        scores = outputs[0].cpu().tolist()
+        scores = outputs[0].squeeze().cpu().tolist()
         return scores
     
     @torch.no_grad()