Update modeling_custom.py

modeling_custom.py

@@ -162,19 +162,19 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
         prompt_embedding = tokens_hidden_states[dummy_iterator, gating_token_positions, :]
         gating_output = self.gating(prompt_embedding)
 
-        #with torch.autocast(device_type=rewards.device.type, dtype=torch.float32):
+        with torch.autocast(device_type=rewards.device.type, dtype=torch.float32):
             # [B, num_quantiles, num_objectives]
-            #reward_quantiles = torch.mul(
-            #    gating_output.unsqueeze(-1).repeat(1, 1, self.num_objectives),
-            #    torch.transpose(rewards, 1, 2)
-            #).sum(1)
+            reward_quantiles = torch.mul(
+                gating_output.unsqueeze(-1).repeat(1, 1, self.num_objectives),
+                torch.transpose(rewards, 1, 2)
+            ).sum(1)
 
         rewards_expectation = rewards.float().mean(dim=2)
 
         score = torch.sum(gating_output.float() * rewards_expectation.float(), dim=1, keepdim=True)
 
         return CustomOutput(
-#            reward_quantiles=reward_quantiles,
+            reward_quantiles=reward_quantiles,
             rewards=rewards_expectation,
             hidden_state=hidden_states,
             prompt_embedding=prompt_embedding,