diff options
| author | junyanz <junyanz@berkeley.edu> | 2017-04-20 01:58:30 -0700 |
|---|---|---|
| committer | junyanz <junyanz@berkeley.edu> | 2017-04-20 01:58:30 -0700 |
| commit | 91e3983d2edd25f61366026e27d0ae2ae9d61282 (patch) | |
| tree | 83d1f6a7803dda081043ea8a8fdcb13577a4c615 /models/networks.py | |
| parent | f576aae4e51c5cb23ec8af58b425f0052912605d (diff) | |
add gpu_ids to save/cuda function
Diffstat (limited to 'models/networks.py')
| -rw-r--r-- | models/networks.py | 9 |
1 files changed, 4 insertions, 5 deletions
diff --git a/models/networks.py b/models/networks.py index 5624a7b..6777359 100644 --- a/models/networks.py +++ b/models/networks.py @@ -32,7 +32,7 @@ def define_G(input_nc, output_nc, ngf, which_model_netG, norm, gpu_ids=[]): if which_model_netG == 'resnet_9blocks': netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, n_blocks=9, gpu_ids=gpu_ids) elif which_model_netG == 'resnet_6blocks': - netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, 6, gpu_ids=gpu_ids) + netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, n_blocks=6, gpu_ids=gpu_ids) elif which_model_netG == 'unet_128': netG = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer, gpu_ids=gpu_ids) elif which_model_netG == 'unet_256': @@ -40,7 +40,7 @@ def define_G(input_nc, output_nc, ngf, which_model_netG, norm, gpu_ids=[]): else: print('Generator model name [%s] is not recognized' % which_model_netG) if len(gpu_ids) > 0: - netG.cuda() + netG.cuda(device_id=gpu_ids[0]) netG.apply(weights_init) return netG @@ -59,7 +59,7 @@ def define_D(input_nc, ndf, which_model_netD, print('Discriminator model name [%s] is not recognized' % which_model_netD) if use_gpu: - netD.cuda() + netD.cuda(device_id=gpu_ids[0]) netD.apply(weights_init) return netD @@ -213,8 +213,7 @@ class UnetGenerator(nn.Module): unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, unet_block) unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, unet_block) unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, unet_block) - unet_block = UnetSkipConnectionBlock(input_nc, ngf, unet_block, - outermost=True) + unet_block = UnetSkipConnectionBlock(output_nc, ngf, unet_block, outermost=True) self.model = unet_block |