1 files changed, 4 insertions, 0 deletions

diff --git a/train_ti.py b/train_ti.py
index 5cbef06..6c74854 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -903,6 +903,10 @@ 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: