[chatgpt]fix train_rm bug with lora (#2741)

2 years ago · 648183a960
1 changed files with 1 additions and 1 deletions
--- a/applications/ChatGPT/chatgpt/trainer/rm.py
+++ b/applications/ChatGPT/chatgpt/trainer/rm.py
@ -43,7 +43,7 @@ class RewardModelTrainer(ABC):
            # train
            if use_lora > 0:
                print("Using Lora")
-                lora.mark_only_lora_as_trainable(self.model)
+                lora.mark_only_lora_as_trainable(self.model.model)
            else:
                self.model.train()
            for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader: