Update train_network.py

This commit is contained in:
gesen2egee
2024-03-11 23:56:09 +08:00
parent 095b8035e6
commit d282c45002

View File

@@ -935,7 +935,7 @@ class NetworkTrainer:
accelerator.end_training() accelerator.end_training()
if is_main_process and args.save_state or args.save_state_on_train_end: if is_main_process and (args.save_state or args.save_state_on_train_end):
train_util.save_state_on_train_end(args, accelerator) train_util.save_state_on_train_end(args, accelerator)
if is_main_process: if is_main_process: