.

ctgan/synthesizers/ctgan.py

@@ -35,19 +35,22 @@ def calc_gradient_penalty(self, real_data, fake_data, device='cpu', pac=10, lamb
         alpha = torch.rand(real_data.size(0) // pac, 1, 1, device=device)
         alpha = alpha.repeat(1, pac, real_data.size(1))
         alpha = alpha.view(-1, real_data.size(1))
+        print(alpha.device)
 
         interpolates = alpha * real_data + ((1 - alpha) * fake_data)
+        a = torch.ones(disc_interpolates.size(), device=device)
+        print(a.device)
 
         disc_interpolates = self(interpolates)
-        self.set_device(device)
         gradients = torch.autograd.grad(
             outputs=disc_interpolates,
             inputs=interpolates,
-            grad_outputs=torch.ones(disc_interpolates.size(), device=device),
+            grad_outputs=a,
             create_graph=True,
             retain_graph=True,
             only_inputs=True,
         )[0]
+        print(gradients.device)
 
         gradients_view = gradients.view(-1, pac * real_data.size(1)).norm(2, dim=1) - 1
         gradient_penalty = ((gradients_view) ** 2).mean() * lambda_