bug Fixes, moved entire dataset to GPU

caffeine-coder1 · caffeine-coder1 · commit db56be58bfcc · 2021-06-30T16:21:31.000+05:30
diff --git a/GAN/WGAN/training.py b/GAN/WGAN/training.py
@@ -46,12 +46,14 @@ def training(opt):
     # ~~~~~~~~~~~~~~~~~~~ loading the dataset ~~~~~~~~~~~~~~~~~~~ #
 
     trans = transforms.Compose(
-        [transforms.ToPILImage(), transforms.Resize((H, W)),
+        [transforms.Resize((H, W)),
          transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
 
     MNIST_data = MNIST('/datasets', True, transform=trans, download=True)
+    MNIST_data = MNIST_data.to(work_device)
 
-    loader = DataLoader(MNIST_data, BATCH_SIZE, True, num_workers=1)
+    loader = DataLoader(MNIST_data, BATCH_SIZE, True,
+                        num_workers=1, pin_memory=True)
 
     # ~~~~~~~~~~~~~~~~~~~ creating tensorboard variables ~~~~~~~~~~~~~~~~~~~ #
 
@@ -95,7 +97,7 @@ def training(opt):
         for batch_idx, (real, _) in enumerate(tqdm(loader)):
             critic.train()
             gen.train()
-            real = real.to(work_device)
+            # real = real.to(work_device)
             fixed_noise = torch.rand(
                 real.shape[0], Z_DIM, 1, 1).to(work_device)
 
