From 91e3983d2edd25f61366026e27d0ae2ae9d61282 Mon Sep 17 00:00:00 2001 From: junyanz Date: Thu, 20 Apr 2017 01:58:30 -0700 Subject: add gpu_ids to save/cuda function --- models/cycle_gan_model.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) (limited to 'models/cycle_gan_model.py') diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py index c3b5b72..eb1c443 100644 --- a/models/cycle_gan_model.py +++ b/models/cycle_gan_model.py @@ -202,11 +202,10 @@ class CycleGANModel(BaseModel): ('real_B', real_B), ('fake_A', fake_A), ('rec_B', rec_B)]) def save(self, label): - use_gpu = self.gpu_ids is not None - self.save_network(self.netG_A, 'G_A', label, use_gpu) - self.save_network(self.netD_A, 'D_A', label, use_gpu) - self.save_network(self.netG_B, 'G_B', label, use_gpu) - self.save_network(self.netD_B, 'D_B', label, use_gpu) + self.save_network(self.netG_A, 'G_A', label, self.gpu_ids) + self.save_network(self.netD_A, 'D_A', label, self.gpu_ids) + self.save_network(self.netG_B, 'G_B', label, self.gpu_ids) + self.save_network(self.netD_B, 'D_B', label, self.gpu_ids) def update_learning_rate(self): lrd = self.opt.lr / self.opt.niter_decay -- cgit v1.2.3-70-g09d2