Cast scores properly

for-ai · Aug 4, 2024 · cca9f33 · cca9f33
1 parent e2091ed
commit cca9f33
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/scripts/run_rewardbench.py b/scripts/run_rewardbench.py
@@ -324,8 +324,9 @@ def main():
             score_rejected_batch = [result["score"] for result in rewards_rejected]
         # for classes that directly output scores (custom code)
         else:
-            score_chosen_batch = rewards_chosen.cpu().numpy().tolist()
-            score_rejected_batch = rewards_rejected.cpu().numpy().tolist()
+            # Cast to float in case of bfloat16
+            score_chosen_batch = rewards_chosen.float().cpu().numpy().tolist()
+            score_rejected_batch = rewards_rejected.float().cpu().numpy().tolist()
 
         # log results
         [