Commit a8c468d6 authored by Guillaume HEUSCH's avatar Guillaume HEUSCH
Browse files

[trainers] fixed the moving of the model to CPU when generating images

parent 4b219490
Pipeline #22459 passed with stage
in 10 minutes and 20 seconds
......@@ -194,9 +194,11 @@ class ConditionalGANTrainer(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.netG = self.netG.cpu()
if self.use_gpu:
self.netG = self.netG.cpu()
fake_examples = self.netG(self.fixed_noise, self.fixed_one_hot)
self.netG = self.netG.cuda()
if self.use_gpu:
self.netG = self.netG.cuda()
vutils.save_image(fake_examples.data, '%s/fake_samples_epoch_%03d.png' % (output_dir, epoch), normalize=True)
# do checkpointing
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment