opt reward model

src/rlhf/reward.py

@@ -68,7 +68,7 @@ class RewardModel(pl.LightningModule):
 
         # reward 得分计算
         self.pred_reward = nn.Sequential(
-            nn.Linear(dim, 1),
+            nn.Linear(dim, 1, bias=False),
             Rearrange('... 1 -> ...')   # 降维
         )
 

train_rm.py

@@ -240,7 +240,10 @@ if __name__ == "__main__":
     rm_model = RewardModel(args)
 
     # 训练
-    trainer = Trainer.from_argparse_args()
+    trainer = Trainer.from_argparse_args(
+        args,
+        callbacks=[train_callback(args)],
+    )
 
     if trainer.global_rank == 0:
         for n in rm_model.state_dict():