fix #175

Former-commit-id: fd557ebb5e3ef2ca330b4d97731af43f4a5a5fc5
2023-07-17 18:07:17 +08:00
parent e9736b2ba0
commit c08ff734a7
5 changed files with 38 additions and 12 deletions
--- a/src/llmtuner/tuner/ppo/trainer.py
+++ b/src/llmtuner/tuner/ppo/trainer.py
@@ -153,7 +153,7 @@ class PPOPeftTrainer(PPOTrainer, PeftTrainer):
            if self.control.should_training_stop:
                break

-    @torch.no_grad()
+    @torch.inference_mode()
    def generate(
        self,
        inputs: Dict[str, torch.Tensor],