From 71a0181fcec3da8768918576b05796be8dddbe0b Mon Sep 17 00:00:00 2001 From: Tong Li Date: Mon, 10 Mar 2025 14:19:10 +0800 Subject: [PATCH] update reward --- .../ColossalChat/coati/distributed/reward/verifiable_reward.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py b/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py index b43ba65c0..ba83f7787 100644 --- a/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py +++ b/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py @@ -21,7 +21,7 @@ class VerifiableReward: # 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: