| 173 | |
| 174 | class _EncoderBlock(nn.Module): |
| 175 | def __init__(self, in_channels, out_channels, downsample=True): |
| 176 | super(_EncoderBlock, self).__init__() |
| 177 | self.downsample = downsample |
| 178 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) |
| 179 | layers = [ |
| 180 | conv3x3(in_channels, out_channels), |
| 181 | nn.BatchNorm2d(out_channels), |
| 182 | nn.ReLU(inplace=True), |
| 183 | conv3x3(out_channels, out_channels), |
| 184 | nn.BatchNorm2d(out_channels), |
| 185 | nn.ReLU(inplace=True), |
| 186 | ] |
| 187 | self.encode = nn.Sequential(*layers) |
| 188 | |
| 189 | def forward(self, x): |
| 190 | if self.downsample: |