Define the output layer
| 397 | |
| 398 | |
| 399 | class FinalBlock2d(nn.Module): |
| 400 | """ |
| 401 | Define the output layer |
| 402 | """ |
| 403 | |
| 404 | def __init__(self, input_nc, output_nc, use_spect=False, tanh_or_sigmoid='tanh'): |
| 405 | super(FinalBlock2d, self).__init__() |
| 406 | |
| 407 | kwargs = {'kernel_size': 7, 'stride': 1, 'padding': 3} |
| 408 | conv = spectral_norm(nn.Conv2d(input_nc, output_nc, **kwargs), use_spect) |
| 409 | |
| 410 | if tanh_or_sigmoid == 'sigmoid': |
| 411 | out_nonlinearity = nn.Sigmoid() |
| 412 | else: |
| 413 | out_nonlinearity = nn.Tanh() |
| 414 | |
| 415 | self.model = nn.Sequential(conv, out_nonlinearity) |
| 416 | |
| 417 | def forward(self, x): |
| 418 | out = self.model(x) |
| 419 | return out |