Fix

1 files changed, 2 insertions, 4 deletions

diff --git a/train_ti.py b/train_ti.py
index 6c74854..97dde1e 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -519,6 +519,8 @@ def main():
     args.seed = args.seed or (torch.random.seed() >> 32)
     set_seed(args.seed)
 
+    save_args(basepath, args)
+
     # Load the tokenizer and add the placeholder token as a additional special token
     if args.tokenizer_name:
         tokenizer = MultiCLIPTokenizer.from_pretrained(args.tokenizer_name)
@@ -903,10 +905,6 @@ def main():
     )
     global_progress_bar.set_description("Total progress")
 
-    save_args(basepath, args, {
-        "global_step": global_step + global_step_offset
-    })
-
     try:
         for epoch in range(num_epochs):
             if accelerator.is_main_process: