diff --git a/src/rlhf/reward.py b/src/rlhf/reward.py
index 665188b6171a899e13932363eb2a32154d9defba..eeafe741f5436c50af04bf534049b8e50029d291 100644
--- a/src/rlhf/reward.py
+++ b/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