Training: Create multiple class images per training image

1 files changed, 4 insertions, 1 deletions

diff --git a/infer.py b/infer.py
index 34e570a..6197aa3 100644
--- a/infer.py
+++ b/infer.py
@@ -262,7 +262,10 @@ def generate(output_dir, pipeline, args):
 
     with autocast("cuda"):
         for i in range(args.batch_num):
-            pipeline.set_progress_bar_config(desc=f"Batch {i + 1} of {args.batch_num}")
+            pipeline.set_progress_bar_config(
+                desc=f"Batch {i + 1} of {args.batch_num}",
+                dynamic_ncols=True
+            )
 
             generator = torch.Generator(device="cuda").manual_seed(seed + i)
             images = pipeline(