Update device

modeling_custom.py

@@ -140,11 +140,11 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
                 # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
                 sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
                 sequence_lengths = sequence_lengths % input_ids.shape[-1]
-                sequence_lengths = sequence_lengths.to("cuda")
+                sequence_lengths = sequence_lengths.to(self.device)
             else:
                 sequence_lengths = -1
 
-        dummy_iterator = torch.arange(batch_size, device=tokens_hidden_states.device)
+        dummy_iterator = torch.arange(batch_size, device=self.device)
         hidden_states = tokens_hidden_states[dummy_iterator, sequence_lengths]
         assert hidden_states.shape == (batch_size, self.config.hidden_size)
         rewards = self.regression_layer(hidden_states)
@@ -163,4 +163,4 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
             gating_output=gating_output,
             score=score,
             logits=score,
-        )
+        )