update webUI, fix #179

Former-commit-id: f9074fed5e22585679661588befcf266a79009f2
2023-07-18 15:35:17 +08:00
parent fd8c2d4aac
commit a864a7b395
9 changed files with 247 additions and 154 deletions
--- a/src/llmtuner/tuner/ppo/trainer.py
+++ b/src/llmtuner/tuner/ppo/trainer.py
@@ -108,7 +108,7 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            replace_model(unwrapped_model, target="reward")
            with torch.no_grad():
                _, _, values = self.model(**self.prepare_model_inputs(queries, responses))
-            rewards = [reward for reward in values[-1].to(torch.float32)] # use float32 type
+            rewards = [reward for reward in values[:, -1].to(torch.float32)] # use float32 type
            replace_model(unwrapped_model, target="default")

            # Run PPO step