Update modeling_custom.py

modeling_custom.py

@@ -195,7 +195,7 @@ class Gemma2ForQuantileSequenceClassification(Gemma2PreTrainedModel):
 
         # [B, num_objectives, num_quantiles, ]
         reward_quantiles = torch.mean(
-            rewards * gating_output.unsqueeze(-1).repeat(1, 1, self.num_quantiles), dim=2)
+            rewards * gating_output.unsqueeze(-1).repeat(1, 1, self.num_quantiles), dim=1)
 
         rewards_expectation = rewards.mean(dim=2)
         score = torch.sum(rewards_expectation.float() * gating_output.float(), dim=-1, keepdim=True)