fix API

Former-commit-id: 9b10c9a12e33ab897056ecc61d977d221c19141b
2023-07-19 00:01:14 +08:00
parent 35e76879f5
commit dc8283d3d7
5 changed files with 12 additions and 9 deletions
--- a/src/llmtuner/tuner/ppo/trainer.py
+++ b/src/llmtuner/tuner/ppo/trainer.py
@@ -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")