save state on train end

2026-04-08 22:35:09 +00:00 · 2024-03-10 23:33:38 +08:00
parent 2d7389185c
commit 095b8035e6
9 changed files with 13 additions and 8 deletions
--- a/sdxl_train.py
+++ b/sdxl_train.py
@@ -712,7 +712,7 @@ def train(args):

    accelerator.end_training()

-    if args.save_state:  # and is_main_process:
+    if args.save_state or args.save_state_on_train_end:        
        train_util.save_state_on_train_end(args, accelerator)

    del accelerator  # この後メモリを使うのでこれは消す