diff options
| author | Taesung Park <taesung_park@berkeley.edu> | 2017-04-20 03:04:48 -0700 |
|---|---|---|
| committer | Taesung Park <taesung_park@berkeley.edu> | 2017-04-20 03:04:48 -0700 |
| commit | 5b177789285aced41c690d5f0d4b412e17693c22 (patch) | |
| tree | 62cbfbd20bedc705df7dc319466622f467553127 /models/base_model.py | |
| parent | 3aeb748074628fd033d0fa798813aa226b8645d6 (diff) | |
| parent | 91e3983d2edd25f61366026e27d0ae2ae9d61282 (diff) | |
Merge branch 'master' of https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
Diffstat (limited to 'models/base_model.py')
| -rw-r--r-- | models/base_model.py | 6 |
1 files changed, 3 insertions, 3 deletions
diff --git a/models/base_model.py b/models/base_model.py index 0ea83d8..ce18635 100644 --- a/models/base_model.py +++ b/models/base_model.py @@ -39,12 +39,12 @@ class BaseModel(): pass # helper saving function that can be used by subclasses - def save_network(self, network, network_label, epoch_label, use_gpu): + def save_network(self, network, network_label, epoch_label, gpu_ids): save_filename = '%s_net_%s.pth' % (epoch_label, network_label) save_path = os.path.join(self.save_dir, save_filename) torch.save(network.cpu().state_dict(), save_path) - if use_gpu and torch.cuda.is_available(): - network.cuda() + if len(gpu_ids) and torch.cuda.is_available(): + network.cuda(device_id=gpu_ids[0]) # helper loading function that can be used by subclasses def load_network(self, network, network_label, epoch_label): |