Update modeling_decision_tree_reward_model.py

modeling_decision_tree_reward_model.py

@@ -93,6 +93,8 @@ class Gemma2ForDecisionTreeRewardModel(Gemma2ForSequenceClassification):
         rewards_2 = embedding_2 @ weight.T + bias
         rewards_diff = rewards_2 - rewards_1
         return {
-            "preference": self.tree.predict(rewards_diff)[0],       "rewards": np.stack([rewards_1, rewards_2]), 
-            "attributes": self.attributes}
+            "preference": self.tree.predict(rewards_diff)[0],
+            "rewards": np.concatenate([rewards_1, rewards_2]),
+            "attributes": self.attributes
+            }