have apex use o1

Author: lucidrains

setup.py

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

stylegan2_pytorch/stylegan2_pytorch.py

@@ -572,7 +572,7 @@ def __init__(self, image_size, latent_dim = 512, fmap_max = 512, style_depth = 8
 
         self.fp16 = fp16
         if fp16:
-            (self.S, self.G, self.D, self.SE, self.GE), (self.G_opt, self.D_opt) = amp.initialize([self.S, self.G, self.D, self.SE, self.GE], [self.G_opt, self.D_opt], opt_level='O2')
+            (self.S, self.G, self.D, self.SE, self.GE), (self.G_opt, self.D_opt) = amp.initialize([self.S, self.G, self.D, self.SE, self.GE], [self.G_opt, self.D_opt], opt_level='O1')
 
     def _init_weights(self):
         for m in self.modules():