| 796 | |
| 797 | class ResBlock(nn.Module): |
| 798 | def __init__(self, in_channels, out_channels=None): |
| 799 | super(ResBlock, self).__init__() |
| 800 | self.in_channels = in_channels |
| 801 | self.out_channels = in_channels if out_channels is None else out_channels |
| 802 | self.norm1 = Normalize(in_channels) |
| 803 | self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| 804 | self.norm2 = Normalize(out_channels) |
| 805 | self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| 806 | if self.in_channels != self.out_channels: |
| 807 | self.conv_out = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) |
| 808 | |
| 809 | def forward(self, x_in): |
| 810 | x = x_in |