diff options
| author | LambdaWill <574819595@qq.com> | 2017-10-21 12:34:44 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2017-10-21 12:34:44 +0800 |
| commit | 1e0fbeaf5bc81c0b9e23655dc31d28c77505b7cd (patch) | |
| tree | d1e066571446be4422fe6b262e2eeb5fb9eba21b /models/networks.py | |
| parent | 9d1bc76e6a4f791a25db1179c7c2b4c62a8d55cd (diff) | |
Compatible with the latest version.
Since the commit `Change device_id to device in python land #3133`, keyword `device_id` has been changed to `device`
Diffstat (limited to 'models/networks.py')
| -rw-r--r-- | models/networks.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/models/networks.py b/models/networks.py index 51e3f25..dca5489 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 netG.cuda(device=gpu_ids[0]) for latest version. init_weights(netD, init_type=init_type) return netD |