diff --git a/wgan_ECAL_HCAL_3crit.py b/wgan_ECAL_HCAL_3crit.py
index d3da95e137d3a0d29200b5337f510fa5539a33ae..01b2f30012979814cf144f2ec6d80f74f1ec7deb 100644
--- a/wgan_ECAL_HCAL_3crit.py
+++ b/wgan_ECAL_HCAL_3crit.py
@@ -242,7 +242,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=True)
-            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)
 
             ## Loss function for ECAL generator