diff --git a/main.py b/main.py
index abedb0a..f852dc2 100755
--- a/main.py
+++ b/main.py
@@ -430,14 +430,16 @@ def train(
             
             images_done += config.train_batch_size
             global_step += 1
-
-            if global_step % 100 == 0:
-                print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r")
-
+            
             if global_step % (config.max_train_steps//20) == 0:
                 progress = (global_step / config.max_train_steps) + 0.05
+                print(f" ---- avg training fps: {images_done / (time.time() - start_time):.2f}", end="\r")
                 yield np.min((progress, 1.0))
 
+            if global_step > config.max_train_steps:
+                print("Reached max steps, stopping training!")
+                break
+
     # final_save
     if (global_step - last_save_step) > 26:
         output_save_dir = f"{checkpoint_dir}/checkpoint-{global_step}"