fix: unit tests

Signed-off-by: Sukriti-Sharma4 <sukriti.sharma4@ibm.com>
foundation-model-stack · Sep 19, 2024 · 146e9f1 · 146e9f1
1 parent 4380c5b
commit 146e9f1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py
@@ -692,7 +692,7 @@ def test_run_causallm_ft_save_with_save_model_dir_save_strategy_no():
         save_model_args.save_strategy = "no"
         save_model_args.output_dir = tempdir
 
-        trainer = sft_trainer.train(MODEL_ARGS, DATA_ARGS, save_model_args, None)
+        trainer, _ = sft_trainer.train(MODEL_ARGS, DATA_ARGS, save_model_args, None)
         logs_path = os.path.join(
             tempdir, FileLoggingTrackerConfig.training_logs_filename
         )