opt reward model

src/rlhf/reward.py

@@ -181,8 +181,10 @@ class RewardModel(pl.LightningModule):
         return reward
     
     def forward(self, x_p, x_a, m_p, m_a):
-        prefer_reward = self.single_forward(x_p, prompt_mask=m_p)
-        alter_reward = self.single_forward(x_a, prompt_mask=m_a)
+        with torch.enable_grad():
+            prefer_reward = self.single_forward(x_p, prompt_mask=m_p)
+        with torch.no_grad():
+            alter_reward = self.single_forward(x_a, prompt_mask=m_a)
 
         return prefer_reward, alter_reward