def __init__(self, channels, use_conv, dims=2, out_channels=None):
    super().__init__()
    self.channels = channels
    self.out_channels = (out_channels or channels)
    self.use_conv = use_conv
    self.dims = dims
    stride = (2 if (dims != 3) else (1, 2, 2))
    if use_conv:
        self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=1)
    else:
        assert (self.channels == self.out_channels)
        self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
def __init__(self, channels, use_conv, dims=2, out_channels=None):
    super().__init__()
    self.channels = channels
    self.out_channels = (out_channels or channels)
    self.use_conv = use_conv
    self.dims = dims
    stride = (2 if (dims != 3) else (1, 2, 2))
    if use_conv:
        self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=1)
    else:
        assert (self.channels == self.out_channels)
        self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)