423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 03:10:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

upcast logits

Browse Source
This commit is contained in:
hiyouga
2024-07-02 22:32:05 +08:00
parent c47ab6c072
commit c13ae2df19
2 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/llamafactory/train/ppo/trainer.py
View File

@@ -407,7 +407,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer):
values = torch.transpose(values, 0, 1)
rewards = values.gather(dim=-1, index=(batch["attention_mask"].sum(dim=-1, keepdim=True) - 1))
return rewards.to(torch.float32).detach().cpu() # use fp32 type
return rewards.float().detach() # use fp32 type
@PPODecorators.empty_device_cache()
def batched_forward_pass(