opt reward model

src/rlhf/reward.py

@@ -119,11 +119,10 @@ class RewardModel(pl.LightningModule):
             state=None,
             extra_embed=extra_embed,
             rm_train=True
-        )
+        )[:, -1, :]
 
         # 所有的 token 向量求平均，并输入到打分模块进行打分
-        pooled = masked_mean(last_token_embeds, mask, dim = 1)
-        reward = self.pred_reward(pooled)
+        reward = self.pred_reward(last_token_embeds)
 
         return reward