code optimization

caffeine-coder1 · caffeine-coder1 · commit b76ca1a3cfff · 2021-07-01T15:10:29.000+05:30
diff --git a/GAN/WGAN/training.py b/GAN/WGAN/training.py
@@ -119,11 +119,17 @@ def training(opt):
                 real.shape[0], Z_DIM, 1, 1).to(work_device)
 
             # ~~~~~~~~~~~~~~~~~~~ critic loop ~~~~~~~~~~~~~~~~~~~ #
-
-            fake = gen(fixed_noise)  # dim of (N,1,W,H)
+            with torch.no_grad():
+                fake = gen(fixed_noise)  # dim of (N,1,W,H)
 
             for _ in range(CRITIC_TRAIN_STEPS):
 
+                critic.zero_grad()
+                # ~~~~~~~~~~~ weight cliping as per WGAN paper ~~~~~~~~~~ #
+
+                for p in critic.parameters():
+                    p.data.clamp_(-WEIGHT_CLIP, WEIGHT_CLIP)
+
                 # ~~~~~~~~~~~~~~~~~~~ forward ~~~~~~~~~~~~~~~~~~~ #
 
                 # make it one dimensional array
@@ -138,16 +144,11 @@ def training(opt):
 
                 # ~~~~~~~~~~~~~~~~~~~ backward ~~~~~~~~~~~~~~~~~~~ #
 
-                critic.zero_grad()
                 C_loss.backward()
                 critic_optim.step()
 
-                # ~~~~~~~~~~~ weight cliping as per WGAN paper ~~~~~~~~~~ #
-
-                for p in critic.parameters():
-                    p.data.clamp_(-WEIGHT_CLIP, WEIGHT_CLIP)
-
             # ~~~~~~~~~~~~~~~~~~~ generator loop ~~~~~~~~~~~~~~~~~~~ #
+            gen.zero_grad()
 
             # ~~~~~~~~~~~~~~~~~~~ forward ~~~~~~~~~~~~~~~~~~~ #
 
@@ -161,7 +162,6 @@ def training(opt):
 
             # ~~~~~~~~~~~~~~~~~~~ backward ~~~~~~~~~~~~~~~~~~~ #
 
-            gen.zero_grad()
             G_loss.backward()
             gen_optim.step()
 
