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/train_textual_inversion_XTI.py
+++ b/train_textual_inversion_XTI.py
@@ -586,7 +586,7 @@ def train(args):

    accelerator.end_training()

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

    updated_embs = text_encoder.get_input_embeddings().weight[token_ids_XTI].data.detach().clone()