add align tensors

modeling_custom.py

@@ -96,6 +96,9 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
                                     temperature=config_dict.get("gating_temperature", 10),
                                     hidden_dim=config_dict.get("gating_hidden_dim", 1024),
                                     n_hidden=config_dict.get("gating_n_hidden", 3))
+    def align_tensor_devices(self, *tensors):
+        target_device = tensors[0].device
+        return [tensor.to(target_device) if tensor.device != target_device else tensor for tensor in tensors]
 
     @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
     def forward(
@@ -153,6 +156,8 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
         prompt_embedding = tokens_hidden_states[dummy_iterator, gating_token_positions, :]
         gating_output = self.gating(prompt_embedding)
 
+        rewards, self.reward_transform_matrix = self.align_tensor_devices(rewards, self.reward_transform_matrix)
+
         rewards_adjusted = rewards @ self.reward_transform_matrix
         score = torch.sum(gating_output * rewards_adjusted, dim=1)