Update glide_text2im/adv.py

glide_text2im/adv.py

@@ -4,7 +4,7 @@ import torch.nn.functional as F
 import torch.optim as optim
 from torch.nn.parallel.distributed import DistributedDataParallel as DDP
 from .nn import mean_flat
-from . import dist_util
+#from . import dist_util
 import functools
 
 class AdversarialLoss(nn.Module):
@@ -16,11 +16,11 @@ class AdversarialLoss(nn.Module):
         self.gan_type = gan_type
         self.gan_k = gan_k
  
-        model = NLayerDiscriminator().to(dist_util.dev()) 
+        model = NLayerDiscriminator().cuda()
         self.discriminator =  DDP(
                 model,
-                device_ids=[dist_util.dev()],
-                output_device=dist_util.dev(),
+                device_ids=[torch.device('cuda')],
+                output_device=torch.device('cuda'),
                 broadcast_buffers=False,
                 bucket_cap_mb=128,
                 find_unused_parameters=False,
@@ -41,7 +41,7 @@ class AdversarialLoss(nn.Module):
             if (self.gan_type.find('WGAN') >= 0):
                 loss_d = (d_fake - d_real).mean()
                 if self.gan_type.find('GP') >= 0:
-                    epsilon = torch.rand(real.size(0), 1, 1, 1).to(dist_util.dev())
+                    epsilon = torch.rand(real.size(0), 1, 1, 1).cuda()
                     epsilon = epsilon.expand(real.size())
                     hat = fake_detach.mul(1 - epsilon) + real.mul(epsilon)
                     hat.requires_grad = True