diff --git a/train.py b/train.py
index b6fe28a..04194f1 100644
--- a/train.py
+++ b/train.py
@@ -267,7 +267,7 @@ def setup_args(args):
         args.lr = args.lr * (total_batch_size**0.55)
         logging.info(f"{Fore.CYAN} * Scaling learning rate {tmp_lr} by {total_batch_size**0.5}, new value: {args.lr}{Style.RESET_ALL}")
 
-    if not os.path.exists(args.save_ckpt_dir):
+    if args.save_ckpt_dir is not None and not os.path.exists(args.save_ckpt_dir):
         os.makedirs(args.save_ckpt_dir)
 
     return args
@@ -410,8 +410,8 @@ def main(args):
 
             clean_prompt = clean_filename(prompt)
 
-            result.save(f"{log_folder}/samples/gs{gs:05}-{clean_prompt[:100]}.jpg", format="JPEG", quality=95, optimize=True, progressive=False)
-            with open(f"{log_folder}/samples/gs{gs:05}-{clean_prompt[:100]}.txt", "w", encoding='utf-8') as f:
+            result.save(f"{log_folder}/samples/gs{gs:05}-{i}-{clean_prompt[:100]}.jpg", format="JPEG", quality=95, optimize=True, progressive=False)
+            with open(f"{log_folder}/samples/gs{gs:05}-{i}-{clean_prompt[:100]}.txt", "w", encoding='utf-8') as f:
                 f.write(prompt)
                 f.write(f"\n seed: {seed}")