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

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


Co-authored-by: lodrick <[email protected]>

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