Move fix to the right place

githubnemo · Jan 24, 2025 · 9b9cf6d · 9b9cf6d
1 parent 8801b5b
commit 9b9cf6d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/test_torch_compile.py b/tests/test_torch_compile.py
@@ -176,6 +176,10 @@ def test_causal_lm_training_trainer_compile(self, settings, tokenizer, data, tmp
             "output_dir": tmp_dir,
             "seed": 0,
         }
+
+        if isinstance(config, AdaLoraConfig):
+            train_kwargs["learning_rate"] = 1e-2
+
         training_args = TrainingArguments(
             torch_compile=not self.fake_compile,
             torch_compile_backend=compile_kwargs.get("torch_compile_backend", None),
@@ -195,7 +199,6 @@ class OptimizerStepCallback(TrainerCallback):
                 def on_optimizer_step(self, args, state, control, **kwargs):
                     model.update_and_allocate(state.global_step)
             trainer.add_callback(OptimizerStepCallback())
-            train_kwargs["learning_rate"] = 1e-2
 
         trainer.train()