|
|
@ -228,7 +228,7 @@ class Downsampling(nn.Module):
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
def __init__(self, in_channels, out_channels,
|
|
|
|
def __init__(self, in_channels, out_channels,
|
|
|
|
kernel_size, stride=1, padding=0,
|
|
|
|
kernel_size, stride=1, padding=0,
|
|
|
|
pre_norm=None, post_norm=None):
|
|
|
|
pre_norm=None, post_norm=None, pre_permute = False):
|
|
|
|
super().__init__()
|
|
|
|
super().__init__()
|
|
|
|
self.pre_norm = pre_norm(in_channels) if pre_norm else nn.Identity()
|
|
|
|
self.pre_norm = pre_norm(in_channels) if pre_norm else nn.Identity()
|
|
|
|
self.pre_permute = pre_permute
|
|
|
|
self.pre_permute = pre_permute
|
|
|
|