added directory creation

Author: eracah

main.py

@@ -5,7 +5,7 @@
 import numpy as np
 from torchvision.utils import save_image
 import time
-
+import os
 
 BATCH_SIZE = 32
 N_EPOCHS = 100
@@ -21,6 +21,10 @@
 
 DEVICE = torch.device('cuda') # torch.device('cpu')
 
+for dir_ in ["models", "samples"]:
+    if not os.path.exists(dir_):
+        os.mkdir(dir_)
+
 preproc_transform = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))