From e1170613a27da597b6362c30885641ae95276f17 Mon Sep 17 00:00:00 2001
From: u010280923 <xinlu.nlp@foxmail.com>
Date: Fri, 10 Mar 2023 15:04:30 +0800
Subject: [PATCH] opt reward model

---
 src/rlhf/reward.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/src/rlhf/reward.py b/src/rlhf/reward.py
index 8e17cea..df77244 100644
--- a/src/rlhf/reward.py
+++ b/src/rlhf/reward.py
@@ -85,12 +85,9 @@ class RewardModel(pl.LightningModule):
         ]
     
     def configure_optimizers(self):
-        # 论文中的参数：lr=1e-5, betas=(0.9, 0.95) 
-        optimizer = torch.optim.Adam([
-            {"rwkv_params": self.rwkv.parameters()},
-            {"rm_params": self.parameters()}
-        ], lr=self.args.lr_init, betas=self.args.betas)
-        
+        # 论文中的参数：
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-5, betas=(0.9, 0.95) )
+        # optimizer = torch.optim.Adam(self.parameters(), lr=self.args.lr_init, betas=self.args.betas)
         
         return optimizer
 
-- 
GitLab