diff --git a/wgan_ECAL_HCAL_3crit.py b/wgan_ECAL_HCAL_3crit.py
index 37a9a9d914fbe80b44fea375cb4ff9135e622c31..0e6bc93fd6dae480faf7cd3e230a815d25549696 100644
--- a/wgan_ECAL_HCAL_3crit.py
+++ b/wgan_ECAL_HCAL_3crit.py
@@ -194,7 +194,7 @@ def train(args, aDE, aDH, aD_H_E, aGE, aGH, device, train_loader, optimizer_d_E,
         fake_ecal = fake_ecal.unsqueeze(1).detach()
 
         z = Variable(Tensor(np.random.uniform(-1, 1, (args.batch_size, args.nz))), requires_grad=False)
-        fake_dataHCAL = aGH(z, label.view(-1, 1, 1, 1, 1), fake_ecal) ## 48 x 30 x 30
+        fake_dataHCAL = aGH(z, label, fake_ecal) ## 48 x 30 x 30
         fake_dataHCAL = fake_dataHCAL.unsqueeze(1).detach()
 
         ## Critic fwd pass on Real