code refactor

caffeine-coder1 · caffeine-coder1 · commit 76f96606fa0a · 2021-07-01T13:04:22.000+05:30
diff --git a/GAN/WGAN/training.py b/GAN/WGAN/training.py
@@ -5,7 +5,7 @@
 from torchvision.transforms import transforms
 from torchvision.datasets import MNIST
 from torch.utils.data import DataLoader
-from model import Discriminator, Faker
+from model import Critic, Faker
 from torch.utils.tensorboard import SummaryWriter
 from pathlib import Path
 import math
@@ -66,8 +66,8 @@ def training(opt):
 
     # ~~~~~~~~~~~~~~~~~~~ loading the model ~~~~~~~~~~~~~~~~~~~ #
 
-    critic = Discriminator(img_channels=CHANNELS,
-                           feature_d=FEATURE_D).to(work_device)
+    critic = Critic(img_channels=CHANNELS,
+                    feature_d=FEATURE_D).to(work_device)
     gen = Faker(Z_DIM, CHANNELS, FEATURE_D).to(work_device)
 
     if opt.resume:
