diff options
| author | Jun-Yan Zhu <junyanz@users.noreply.github.com> | 2017-11-02 15:30:12 -0700 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2017-11-02 15:30:12 -0700 |
| commit | a303e4b0c174c8eb86e53f79b267df2e344da81a (patch) | |
| tree | cc0f69f9fc287133d511a265e1fd95a63fcd8f08 /models | |
| parent | 822896085f183859c41fd8f5fc6e9dbbfb4c706c (diff) | |
| parent | 343dda259b4b6a64a338fa61a9f1c70893b8fc7e (diff) | |
Merge pull request #130 from LambdaWill/patch-1
Compatible with the latest version.
Diffstat (limited to 'models')
| -rw-r--r-- | models/base_model.py | 2 | ||||
| -rw-r--r-- | models/networks.py | 4 |
2 files changed, 3 insertions, 3 deletions
diff --git a/models/base_model.py b/models/base_model.py index 446a903..d62d189 100644 --- a/models/base_model.py +++ b/models/base_model.py @@ -44,7 +44,7 @@ class BaseModel(): save_path = os.path.join(self.save_dir, save_filename) torch.save(network.cpu().state_dict(), save_path) if len(gpu_ids) and torch.cuda.is_available(): - network.cuda(device_id=gpu_ids[0]) + network.cuda(device_id=gpu_ids[0]) # network.cuda(device=gpu_ids[0]) for the latest version. # helper loading function that can be used by subclasses def load_network(self, network, network_label, epoch_label): diff --git a/models/networks.py b/models/networks.py index 51e3f25..949659d 100644 --- a/models/networks.py +++ b/models/networks.py @@ -118,7 +118,7 @@ def define_G(input_nc, output_nc, ngf, which_model_netG, norm='batch', use_dropo else: raise NotImplementedError('Generator model name [%s] is not recognized' % which_model_netG) if len(gpu_ids) > 0: - netG.cuda(device_id=gpu_ids[0]) + netG.cuda(device_id=gpu_ids[0]) # or netG.cuda(device=gpu_ids[0]) for latest version. init_weights(netG, init_type=init_type) return netG @@ -139,7 +139,7 @@ def define_D(input_nc, ndf, which_model_netD, raise NotImplementedError('Discriminator model name [%s] is not recognized' % which_model_netD) if use_gpu: - netD.cuda(device_id=gpu_ids[0]) + netD.cuda(device_id=gpu_ids[0]) # or netD.cuda(device=gpu_ids[0]) for latest version. init_weights(netD, init_type=init_type) return netD |