Misc fixes

This commit is contained in:
James Betker 2021-01-04 10:53:53 -07:00
parent 4d8064c32c
commit 5916f5f7d4
2 changed files with 2 additions and 1 deletions

View File

@ -168,7 +168,7 @@ class Discriminator_VGG_128_GN(nn.Module):
@register_model
def register_discriminator_vgg_128(opt_net, opt):
def register_discriminator_vgg_128_gn(opt_net, opt):
return Discriminator_VGG_128_GN(in_nc=opt_net['in_nc'], nf=opt_net['nf'],
input_img_factor=opt_net['image_size'] / 128,
extra_conv=opt_get(opt_net, ['extra_conv'], False),

View File

@ -811,6 +811,7 @@ class StyleGan2Discriminator(nn.Module):
self.flatten = Flatten()
if mlp:
self.to_logit = nn.Sequential(nn.Linear(latent_dim, 100),
leaky_relu(),
nn.Linear(100, 1))
else:
self.to_logit = nn.Linear(latent_dim, 1)