From c49b3bfd7e76dc7af4b849214c31292459ec26b4 Mon Sep 17 00:00:00 2001 From: Harikrishnan Balagopal Date: Thu, 25 Apr 2024 14:12:33 +0530 Subject: [PATCH] fix: bug with AIM server availability check Signed-off-by: Harikrishnan Balagopal --- tuning/aim_loader.py | 2 +- tuning/sft_trainer.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tuning/aim_loader.py b/tuning/aim_loader.py index 32a2a9f18..a74bcc627 100644 --- a/tuning/aim_loader.py +++ b/tuning/aim_loader.py @@ -38,6 +38,6 @@ def get_aimstack_callback(): if aim_db: aim_callback = AimCallback(repo=aim_db, experiment=aim_experiment) else: - aim_callback = AimCallback(experiment=aim_experiment) + aim_callback = None return aim_callback diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py index 1f8effa28..0f55196a0 100644 --- a/tuning/sft_trainer.py +++ b/tuning/sft_trainer.py @@ -244,7 +244,9 @@ def train( callbacks = [FileLoggingCallback(logger)] if is_aim_available(): - callbacks.append(get_aimstack_callback()) + aimstack_callback = get_aimstack_callback() + if aimstack_callback is not None: + callbacks.append(aimstack_callback) if (trainer_controller_args is not None) and ( trainer_controller_args.trainer_controller_config_file is not None