Update modeling_custom.py

modeling_custom.py

@@ -166,7 +166,7 @@ class LlamaForRewardModelWithGating(LlamaPreTrainedModel):
         with torch.autocast(device_type=rewards.device.type, dtype=torch.float32):
             # [B, num_quantiles, num_objectives]
             reward_quantiles = torch.mean(
-                gating_output.unsqueeze(-1).repeat(1, 1, self.num_objectives) * rewards,
+                gating_output.unsqueeze(-1).repeat(1, 1, self.num_quantiles) * rewards,
                 dim=1
             )