move tensor if necessary (like device_map: auto with >1 gpu)

modeling_internlm2.py

@@ -1880,6 +1880,10 @@ class InternLM2ForRewardModel(InternLM2PreTrainedModel):
         # get end reward token's score
         ends = attention_mask.cumsum(dim=1).argmax(dim=1).view(-1,1)
 
+        # move tensor if on different device so auto_map works
+        if ends.device != hidden_states.device:
+            ends = ends.to(hidden_states.device)
+
         reward_scores = torch.gather(hidden_states.squeeze(-1), 1, ends)
 
         loss = None