Update modeling_custom.py

modeling_custom.py

@@ -160,7 +160,7 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
 
         gating_token_positions = [find_token_for_gating(ids.tolist()) for ids in input_ids]
         prompt_embedding = tokens_hidden_states[dummy_iterator, gating_token_positions, :]
-        gating_output = self.gating(prompt_embedding.float())
+        gating_output = self.gating(prompt_embedding)
 
         #with torch.autocast(device_type=rewards.device.type, dtype=torch.float32):
             # [B, num_quantiles, num_objectives]
@@ -171,7 +171,7 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
 
         rewards_expectation = rewards.float().mean(dim=2)
 
-        score = torch.sum(gating_output.float() * rewards_expectation, dim=1, keepdim=True)
+        score = torch.sum(gating_output.float() * rewards_expectation.float(), dim=1, keepdim=True)
 
         return CustomOutput(
 #            reward_quantiles=reward_quantiles,