no sync all DDPs that were usedeven if detached, to play it safe

Author: lucidrains

setup.py

@@ -8,7 +8,7 @@
           'stylegan2_pytorch = stylegan2_pytorch.cli:main',
       ],
   },
-  version = '0.22.1',
+  version = '0.22.2',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

stylegan2_pytorch/stylegan2_pytorch.py

@@ -829,7 +829,7 @@ def train(self):
         avg_pl_length = self.pl_mean
         self.GAN.D_opt.zero_grad()
 
-        for i in gradient_accumulate_contexts(self.gradient_accumulate_every, self.is_ddp, ddps=[D_aug]):
+        for i in gradient_accumulate_contexts(self.gradient_accumulate_every, self.is_ddp, ddps=[D_aug, S, G]):
             get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list
             style = get_latents_fn(batch_size, num_layers, latent_dim, device=self.rank)
             noise = image_noise(batch_size, image_size, device=self.rank)