From 0c942106bf9f2dbb9ee533ac00f302b34dda2d9d Mon Sep 17 00:00:00 2001 From: Evgeny Chukreev Date: Thu, 18 May 2023 20:09:06 +0200 Subject: [PATCH] fix: don't save state if no --save-state arg given --- library/train_util.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/library/train_util.py b/library/train_util.py index 0bd87bc8..5d68c6ac 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -3167,10 +3167,11 @@ def save_sd_model_on_epoch_end_or_stepwise( print(f"removing old model: {remove_out_dir}") shutil.rmtree(remove_out_dir) - if on_epoch_end: - save_and_remove_state_on_epoch_end(args, accelerator, epoch_no) - else: - save_and_remove_state_stepwise(args, accelerator, global_step) + if args.save_state: + if on_epoch_end: + save_and_remove_state_on_epoch_end(args, accelerator, epoch_no) + else: + save_and_remove_state_stepwise(args, accelerator, global_step) def save_and_remove_state_on_epoch_end(args: argparse.Namespace, accelerator, epoch_no):