Fix issues with generate

isaacsquires · Feb 7, 2022 · defaeda · defaeda
1 parent 2f0a713
commit defaeda
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/config.py b/config.py
@@ -22,6 +22,7 @@ def __init__(self, tag):
         self.Lambda = 10
         self.critic_iters = 10
         self.lz = 4
+        self.lf = 4
         self.ngpu = 1
         if self.ngpu > 0:
             self.device_name = "cuda:0"

diff --git a/main.py b/main.py
@@ -19,18 +19,18 @@ def main(mode, offline, tag):
     :raises ValueError: [description]
     """
     print("Running in {} mode, tagged {}, offline {}".format(mode, tag, offline))
-    overwrite = util.check_existence(tag)
-    util.initialise_folders(tag, overwrite)
 
     # Initialise Config object
     c = Config(tag)
 
     if mode == 'train':
+        overwrite = util.check_existence(tag)
+        util.initialise_folders(tag, overwrite)
         netD, netG = networks.make_nets(c, overwrite)
         train(c, netG, netD, offline=offline, overwrite=overwrite)
 
     elif mode == 'generate':
-        netD, netG = networks.make_nets(c, Training=0)
+        netD, netG = networks.make_nets(c, training=0)
         net_g = netG()
         util.generate(c, net_g)
         print("Img generated")

diff --git a/src/util.py b/src/util.py
@@ -220,7 +220,7 @@ def generate(c, netG):
         netG = nn.DataParallel(netG, list(range(ngpu))).to(device)
     netG.load_state_dict(torch.load(f"{pth}/Gen.pt"))
     netG.eval()
-    noise = torch.randn(1, nz, lf, lf, lf)
+    noise = torch.randn(1, nz, lf, lf)
     raw = netG(noise)
     gb = post_process(raw)
     tif = np.array(gb[0], dtype=np.uint8)