refactor: update test_run_with_additional_callbacks() unit-test

Signed-off-by: Vassilis Vassiliadis <[email protected]>
foundation-model-stack · Jun 26, 2024 · e56ee75 · e56ee75
1 parent 6c966a7
commit e56ee75
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -629,8 +629,7 @@ def test_run_with_additional_callbacks():
     """Ensure that train() can work with additional_callbacks"""
 
     with tempfile.TemporaryDirectory() as tempdir:
-        train_args = copy.deepcopy(TRAIN_ARGS)
-        train_args.output_dir = tempdir
+        train_args = get_train_args(output_dir=tempdir)
         model_args = copy.deepcopy(MODEL_ARGS)
 
         sft_trainer.train(