github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-09 21:09:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix bug, tested

Browse Source
This commit is contained in:
YeAnbang
2025-06-09 09:37:28 +08:00
parent dc3033e68a
commit 3bed6ae9ee
3 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
applications/ColossalChat/coati/distributed/producer.py
View File

@@ -291,7 +291,7 @@ class BaseProducer:
reward_model_output = self.reward_model(
outputs["input_ids"].view((-1, outputs["input_ids"].size(-1))),
gt_answer=gt_answer,
response_idx=outputs["response_idx"],
response_idx=outputs["response_idx"].view((-1, 2)),
)
outputs["reward"] = (
torch.tensor([value[0] for value in reward_model_output])