wait a bit before doing path length regularization and gradient penalty

Author: lucidrains

setup.py

@@ -4,7 +4,7 @@
   name = 'stylegan2_pytorch',
   packages = find_packages(),
   scripts=['bin/stylegan2_pytorch'],
-  version = '0.18.0',
+  version = '0.18.1',
   license='GPLv3+',
   description = 'StyleGan2 in Pytorch',
   author = 'Phil Wang',

stylegan2_pytorch/stylegan2_pytorch.py

@@ -699,8 +699,8 @@ def train(self):
 
         aug_prob   = self.aug_prob
 
-        apply_gradient_penalty = self.steps % 4 == 0
-        apply_path_penalty = self.steps % 32 == 0
+        apply_gradient_penalty = self.steps > 5000 and self.steps % 4 == 0
+        apply_path_penalty = self.steps > 5000 and self.steps % 32 == 0
         apply_cl_reg_to_generated = self.steps > 20000
 
         backwards = partial(loss_backwards, self.fp16)