Define the output layer
| 348 | return out |
| 349 | |
| 350 | class FinalBlock2d(nn.Module): |
| 351 | """ |
| 352 | Define the output layer |
| 353 | """ |
| 354 | def __init__(self, input_nc, output_nc, use_spect=False, tanh_or_sigmoid='tanh'): |
| 355 | super(FinalBlock2d, self).__init__() |
| 356 | |
| 357 | kwargs = {'kernel_size': 7, 'stride': 1, 'padding':3} |
| 358 | conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect) |
| 359 | |
| 360 | if tanh_or_sigmoid == 'sigmoid': |
| 361 | out_nonlinearity = nn.Sigmoid() |
| 362 | else: |
| 363 | out_nonlinearity = nn.Tanh() |
| 364 | |
| 365 | self.model = nn.Sequential(conv, out_nonlinearity) |
| 366 | def forward(self, x): |
| 367 | out = self.model(x) |
| 368 | return out |