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

fix API

Browse Source
This commit is contained in:
hiyouga
2023-07-19 00:01:14 +08:00
parent fe2887ca13
commit 29af67b015
5 changed files with 12 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/llmtuner/tuner/ppo/trainer.py
View File

@@ -107,7 +107,11 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
# Compute rewards
replace_model(unwrapped_model, target="reward")
with torch.no_grad():
_, _, values = self.model(**self.prepare_model_inputs(queries, responses))
_, _, values = self.model(
**self.prepare_model_inputs(queries, responses),
output_hidden_states=True,
return_dict=True
)
rewards = [reward for reward in values[:, -1].to(torch.float32)] # use float32 type
replace_model(unwrapped_model, target="default")