.

sdv-dev · Jul 29, 2024 · 3ad1f48 · 3ad1f48
1 parent 118f3b6
commit 3ad1f48
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ctgan/synthesizers/ctgan.py b/ctgan/synthesizers/ctgan.py
@@ -39,7 +39,7 @@ def calc_gradient_penalty(self, real_data, fake_data, device='cpu', pac=10, lamb
         interpolates = alpha * real_data + ((1 - alpha) * fake_data)
 
         disc_interpolates = self(interpolates)
-        self.set_device(device)
+        torch.cuda.set_device(0)
         gradients = torch.autograd.grad(
             outputs=disc_interpolates,
             inputs=interpolates,