prevent circular imports

axolotl-ai-cloud · Feb 12, 2024 · 4b3c1e7 · 4b3c1e7
1 parent 8b08d4a
commit 4b3c1e7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/axolotl/monkeypatch/relora.py b/src/axolotl/monkeypatch/relora.py
@@ -24,7 +24,6 @@
 )
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
 
-from axolotl.core.trainer_builder import AxolotlTrainingArguments
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import barrier, is_main_process
 
@@ -114,7 +113,7 @@ def on_train_begin(
 
     def on_step_begin(
         self,
-        args: AxolotlTrainingArguments,
+        args: TrainingArguments,
         state: TrainerState,
         control: TrainerControl,
         model: peft.LoraModel,