Skip to content

GitLab

Unverified Commit 648183a9 authored by BlueRum's avatar BlueRum Committed by GitHub
Browse files

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

parent b6e3b955
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
applications/ChatGPT/chatgpt/trainer/rm.py
View file @ 648183a9
......@@ -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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment