bug fixes for loading data

Andrewzh112 · Dec 11, 2020 · 42c7832 · 42c7832
1 parent be17ed8
commit 42c7832
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/sagan/train.py b/sagan/train.py
@@ -17,8 +17,8 @@
 parser = argparse.ArgumentParser()
 
 # model parameters
-parser.add_argument('--h_dim', type=float, default=64, help='model dimensions multiplier')
-parser.add_argument('--z_dim', type=float, default=100, help='dimension of random noise latent vector')
+parser.add_argument('--h_dim', type=int, default=64, help='model dimensions multiplier')
+parser.add_argument('--z_dim', type=int, default=100, help='dimension of random noise latent vector')
 
 # data paramters
 parser.add_argument('--img_size', type=int, default=128, help='H, W of the input images')
@@ -65,7 +65,7 @@ def train():
     optimizer_D = torch.optim.Adam(D.parameters(), lr=opt.lr_D, betas=opt.betas)
 
     loader = get_loaders(opt.data_path, opt.img_ext, opt.crop_size,
-                         opt.img_size, opt.batch_size)
+                         opt.img_size, opt.batch_size, opt.download)
     # sample fixed z to see progress through training
     fixed_z = torch.randn(opt.sample_size, opt.z_dim).to(device)
 

diff --git a/vae/main.py b/vae/main.py
@@ -19,6 +19,7 @@
 parser.add_argument('--img_channels', type=int, default=3, help='Numer of channels for images')
 parser.add_argument('--img_size', type=int, default=64, help='H, W of the input images')
 parser.add_argument('--crop_size', type=int, default=128, help='H, W of the input images')
+parser.add_argument('--download', action="store_true", default=False, help='If auto download CelebA dataset')
 
 # model params
 parser.add_argument('--z_dim', type=float, default=100, help='dimension of random noise latent vector')