opt reward model

src/rlhf/reward.py

@@ -71,10 +71,12 @@ class RewardModel(pl.LightningModule):
         self.padding_embed = nn.Parameter(torch.zeros(dim), requires_grad=False)
 
         # reward 得分计算
-        self.pred_reward = nn.Sequential(
-            nn.Linear(dim, 1, bias=False),
-            Rearrange('... 1 -> ...')   # 降维
-        )
+        self.pred_reward = nn.Linear(dim, 1, bias=False)
+
+        # self.pred_reward = nn.Sequential(
+        #     nn.Linear(dim, 1, bias=False),
+        #     Rearrange('... 1 -> ...')   # 降维
+        # )
 
     def load(self, path):
         path = Path(path)
@@ -177,6 +179,7 @@ class RewardModel(pl.LightningModule):
 
         # 计算奖励
         reward = self.pred_reward(last_token_embeds)
+        reward = reward.squeeze(-1)
 
         return reward