From fe5358bacc0298fbb69e33738621ee812cc0ada6 Mon Sep 17 00:00:00 2001
From: Erland366 <erland.pg366@gmail.com>
Date: Thu, 3 Jul 2025 15:39:56 +0000
Subject: [PATCH] Refactor UnslothTrainingArguments to support fallback for
 TrainingArguments import

---
 unsloth/trainer.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/unsloth/trainer.py b/unsloth/trainer.py
index a360205432..e7c73deeaf 100644
--- a/unsloth/trainer.py
+++ b/unsloth/trainer.py
@@ -61,9 +61,13 @@ else:
     pass
 pass
 
-from trl import SFTConfig
-@dataclass
-class UnslothTrainingArguments(SFTConfig):
+try:
+    from trl import SFTConfig as TrainingArguments
+except:
+    from transformers import TrainingArguments
+pass
+
+class UnslothTrainingArguments(TrainingArguments):
     embedding_learning_rate : Optional[float] = field(
         default = None,
         metadata = {"help" : "Different learning rates for embeddings and lm_head."}