2 files changed, 8 insertions, 0 deletions

diff --git a/train_dreambooth.py b/train_dreambooth.py
index 05f6cb5..cd0bf67 100644
--- a/train_dreambooth.py
+++ b/train_dreambooth.py
@@ -955,6 +955,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:
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: