From 6386de4cced6d3f731e73c6e64ac3ebb3008988d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 5 Feb 2025 05:36:51 -0800
Subject: [PATCH] Update rl.py

---
 unsloth/models/rl.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index f8d4d5412c..40979ec76f 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -151,15 +151,16 @@ def _PatchRLStatistics(metrics):
 pass
 
 
-def PatchRLStatistics(algorithm = "grpo"):
-    if algorithm == "grpo":
+def PatchRLStatistics(algorithm = "GRPO"):
+    algorithm = algorithm.upper()
+    if algorithm == "GRPO":
         metrics = [
             "completion_length",
             "reward",
             "reward_std",
             "kl",
         ]
-    elif algorithm == "dpo" or algorithm == "kto":
+    elif algorithm == "DPO" or algorithm == "KTO":
         metrics = [
             "rewards/chosen",
             "rewards/rejected",