do not break previous models

Author: lucidrains

setup.py

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

stylegan2_pytorch/stylegan2_pytorch.py

@@ -977,6 +977,11 @@ def load(self, num = -1):
         self.steps = name * self.save_every
 
         load_data = torch.load(self.model_name(name))
+
+        # make backwards compatible
+        if 'GAN' not in load_data:
+            load_data = {'GAN': load_data}
+
         self.GAN.load_state_dict(load_data['GAN'])
 
         if self.GAN.fp16 and 'amp' in load_data: