diff --git a/main.py b/main.py
index 9e5cdd9..eddf511 100644
--- a/main.py
+++ b/main.py
@@ -96,7 +96,11 @@
     if args.mode == "test":
         out_dir = os.path.dirname(args.pretrain_path)
         out_dir_img = os.path.join(out_dir, "images", "test")
-        #os.mkdir(out_dir_img)
+        
+        try:
+            os.mkdir(out_dir_img)
+        except:
+            continue
 
         # load data
         test_loader = dataloader.get_dataloader(os.path.join(args.data_dir, "testA"),