add the option to work in CPU

.gitignore

@@ -5,4 +5,5 @@ Lib/*
 logs/*
 WiggleGAN_mod.py
 WiggleGAN_noCycle.py
-pyvenv.cfg
+pyvenv.cfg
+py

WiggleGAN.py

@@ -783,9 +783,13 @@ class WiggleGAN(object):
     def load(self):
         save_dir = os.path.join(self.save_dir, self.dataset, self.model_name)
 
-        self.G.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_' + self.seed_load + '_G.pkl')))
+        map_loc=None
+        if not torch.cuda.is_available():
+            map_loc='cpu'
+
+        self.G.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_' + self.seed_load + '_G.pkl'), map_location=map_loc))
         if not self.wiggle:
-            self.D.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_' + self.seed_load + '_D.pkl')))
+            self.D.load_state_dict(torch.load(os.path.join(save_dir, self.model_name + '_' + self.seed_load + '_D.pkl'), map_location=map_loc))
 
     def wiggleEf(self):
         seed, epoch = self.seed_load.split('_')