update reward

Tong Li · Tong Li · commit 71a0181fcec3 · 2025-03-10T14:19:10.000+08:00
diff --git a/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py b/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py
@@ -21,7 +21,7 @@ def __call__(
         # Get batch size
         bs = input_ids.size(0)
         # Initialize reward
-        rewards = torch.zeros(bs, device=input_ids.device)
+        rewards = torch.zeros((bs, 3), device=input_ids.device)
 
         # Loop through reward functions
         for reward_fn in self.reward_fns:
