[trainer] removed the movement of networks to CPU/GPU when saving images ...

bob/learn/pytorch/trainers/DRGANTrainer.py

@@ -132,6 +132,11 @@ class DRGANTrainer(object):
       fixed_one_hot[k, k] = 1
     fixed_one_hot = Variable(fixed_one_hot)
 
+    if self.use_gpu:
+      fixed_image = fixed_image.cuda()
+      fixed_noise = fixed_noise.cuda()
+      fixed_one_hot = fixed_one_hot.cuda()
+
     return fixed_image, fixed_noise, fixed_one_hot
    
 
@@ -450,8 +455,8 @@ class DRGANTrainer(object):
       # tried to move tensors, variables on the GPU -> does not work
       # let the tensors on the CPU -> does not work
       # => model has to be brought back to the CPU :/
-      self.encoder = self.encoder.cpu()
-      self.decoder = self.decoder.cpu()
+      #self.encoder = self.encoder.cpu()
+      #self.decoder = self.decoder.cpu()
       
       fixed_encoded_id = self.encoder(fixed_image)
       fake_examples = self.decoder(fixed_noise, fixed_one_hot, fixed_encoded_id)
@@ -460,9 +465,9 @@ class DRGANTrainer(object):
       to_save[1:] = fake_examples.data
       vutils.save_image(to_save, '%s/fake_samples_epoch_%03d.png' % (images_dir, epoch), normalize=True)
       
-      if torch.cuda.is_available():
-        self.encoder = self.encoder.cuda()
-        self.decoder = self.decoder.cuda()
+      #if torch.cuda.is_available():
+      #  self.encoder = self.encoder.cuda()
+      #  self.decoder = self.decoder.cuda()
 
       # do checkpointing
       torch.save(self.encoder.state_dict(), '%s/encoder_epoch_%d.pth' % (models_dir, epoch))
@@ -474,7 +479,7 @@ class DRGANTrainer(object):
       model_files = glob.glob(models_dir + '/*.pth') 
       for model_file in model_files:
         model_epoch = int(model_file.split('_')[-1].split('.')[0])
-        if ((model_epoch % keep_model) != 0) and (model_epoch != epoch):
+        if (((model_epoch % keep_model) != 0) and (model_epoch != epoch)):
           os.remove(model_file)
           logger.info("{} removed !".format(model_file))