def __init__(self, in_channels, out_channels, internal_ratio=4, kernel_size=2, padding=0, dropout_prob=0.0, bias=False, non_linear='ReLU'):
    super().__init__()
    internal_channels = (in_channels // internal_ratio)
    self.main_conv1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(out_channels))
    self.ext_conv1 = nn.Sequential(nn.Conv2d(in_channels, internal_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(internal_channels), NON_LINEARITY[non_linear])
    self.ext_conv2 = nn.Sequential(nn.ConvTranspose2d(internal_channels, internal_channels, kernel_size=kernel_size, stride=2, padding=padding, output_padding=0, bias=bias), nn.BatchNorm2d(internal_channels), NON_LINEARITY[non_linear])
    self.ext_conv3 = nn.Sequential(nn.Conv2d(internal_channels, out_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(out_channels), NON_LINEARITY[non_linear])
    self.ext_regul = nn.Dropout2d(p=dropout_prob)
    self.out_prelu = NON_LINEARITY[non_linear]
def __init__(self, in_channels, out_channels, internal_ratio=4, kernel_size=2, padding=0, dropout_prob=0.0, bias=False, non_linear='ReLU'):
    super().__init__()
    internal_channels = (in_channels // internal_ratio)
    self.main_conv1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(out_channels))
    self.ext_conv1 = nn.Sequential(nn.Conv2d(in_channels, internal_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(internal_channels), NON_LINEARITY[non_linear])
    self.ext_conv2 = nn.Sequential(nn.ConvTranspose2d(internal_channels, internal_channels, kernel_size=kernel_size, stride=2, padding=padding, output_padding=0, bias=bias), nn.BatchNorm2d(internal_channels), NON_LINEARITY[non_linear])
    self.ext_conv3 = nn.Sequential(nn.Conv2d(internal_channels, out_channels, kernel_size=1, bias=bias), nn.BatchNorm2d(out_channels), NON_LINEARITY[non_linear])
    self.ext_regul = nn.Dropout2d(p=dropout_prob)
    self.out_prelu = NON_LINEARITY[non_linear]