Update modeling_custom.py

modeling_custom.py

@@ -140,7 +140,7 @@ 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(tokens_hidden_states.device)
+                sequence_lengths = sequence_lengths.to("cuda")
             else:
                 sequence_lengths = -1