jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, gin_channels=0):
        super(iSTFT_Generator, self).__init__()
        # self.h = h
        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size

        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = weight_norm(
            Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u)//2)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.post_n_fft = self.gen_istft_n_fft
        self.conv_post = weight_norm(
            Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
        self.ups.apply(commons.init_weights)
        self.conv_post.apply(commons.init_weights)
        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
        self.cond = nn.Conv1d(256, 512, 1)
        self.stft = stft_module.TorchSTFT(filter_length=self.gen_istft_n_fft,
                                          hop_length=self.gen_istft_hop_size, win_length=self.gen_istft_n_fft)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, gin_channels=0):
        super(iSTFT_Generator, self).__init__()
        # self.h = h
        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size

        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(torch.nn.utils.parametrizations.weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u)//2)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.post_n_fft = self.gen_istft_n_fft
        self.conv_post = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(ch, self.post_n_fft + 2, 7, 1, padding=3))
        self.ups.apply(commons.init_weights)
        self.conv_post.apply(commons.init_weights)
        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
        self.cond = nn.Conv1d(256, 512, 1)
        self.stft = stft_module.TorchSTFT(filter_length=self.gen_istft_n_fft,
                                          hop_length=self.gen_istft_hop_size, win_length=self.gen_istft_n_fft)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, subbands, gin_channels=0):
        super(Multiband_iSTFT_Generator, self).__init__()
        # self.h = h
        self.subbands = subbands
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = weight_norm(
            Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u+1-i)//2, output_padding=1-i)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.post_n_fft = gen_istft_n_fft
        self.ups.apply(commons.init_weights)
        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
        self.reshape_pixelshuffle = []

        self.subband_conv_post = weight_norm(
            Conv1d(ch, self.subbands*(self.post_n_fft + 2), 7, 1, padding=3))

        self.subband_conv_post.apply(commons.init_weights)
        self.cond = nn.Conv1d(256, 512, 1)
        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, subbands, gin_channels=0):
        super(Multiband_iSTFT_Generator, self).__init__()
        # self.h = h
        self.subbands = subbands
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(torch.nn.utils.parametrizations.weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u+1-i)//2, output_padding=1-i)))

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.post_n_fft = gen_istft_n_fft
        self.ups.apply(commons.init_weights)
        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
        self.reshape_pixelshuffle = []

        self.subband_conv_post = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(ch, self.subbands*(self.post_n_fft + 2), 7, 1, padding=3))

        self.subband_conv_post.apply(commons.init_weights)
        self.cond = nn.Conv1d(256, 512, 1)
        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d):
                torch.nn.utils.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d):
                torch.nn.utils.parametrizations.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def weight_norm(self):
        self.depth_conv = weight_norm(self.depth_conv, name="weight")
        self.point_conv = weight_norm(self.point_conv, name="weight")
//  [python]
def weight_norm(self):
        self.depth_conv = torch.nn.utils.parametrizations.weight_norm(self.depth_conv, name="weight")
        self.point_conv = torch.nn.utils.parametrizations.weight_norm(self.point_conv, name="weight")
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
                torch.nn.utils.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def apply_weight_norm(self):
        """Apply weight normalization module from all of the layers."""

        def _apply_weight_norm(m):
            if isinstance(m, nn.Conv1d) or isinstance(m, nn.ConvTranspose1d):
                torch.nn.utils.parametrizations.weight_norm(m)

        self.apply(_apply_weight_norm)
//  [python]
def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size):
        super().__init__()
        assert kernel_size % 2 == 1
        assert n_channels % 2 == 0
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.upsample = torch.nn.Upsample(scale_factor=2, mode='nearest')

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.weight_norm(start, name='weight')
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        cond_layer = torch.nn.Conv1d(n_mel_channels, 2 * n_channels * n_layers, 1)
        self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        padding = (kernel_size - 1) // 2
        for _ in range(n_layers):
            self.in_layers.append(
                torch.nn.Sequential(
                    torch.nn.BatchNorm1d(n_channels),
                    torch.nn.Conv1d(n_channels, n_channels, kernel_size, padding=padding, groups=n_channels),
                    torch.nn.Conv1d(n_channels, 2 * n_channels, 1),
                )
            )

            res_skip_layer = torch.nn.Conv1d(n_channels, n_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, kernel_size):
        super().__init__()
        assert kernel_size % 2 == 1
        assert n_channels % 2 == 0
        self.n_layers = n_layers
        self.n_channels = n_channels
        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.upsample = torch.nn.Upsample(scale_factor=2, mode='nearest')

        start = torch.nn.Conv1d(n_in_channels, n_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start, name='weight')
        self.start = start

        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  This helps with training stability
        end = torch.nn.Conv1d(n_channels, 2 * n_in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        cond_layer = torch.nn.Conv1d(n_mel_channels, 2 * n_channels * n_layers, 1)
        self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

        padding = (kernel_size - 1) // 2
        for _ in range(n_layers):
            self.in_layers.append(
                torch.nn.Sequential(
                    torch.nn.BatchNorm1d(n_channels),
                    torch.nn.Conv1d(n_channels, n_channels, kernel_size, padding=padding, groups=n_channels),
                    torch.nn.Conv1d(n_channels, 2 * n_channels, 1),
                )
            )

            res_skip_layer = torch.nn.Conv1d(n_channels, n_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def weight_norm_modules(module, name="weight", dim=0):
    if isinstance(module, Depthwise_Separable_Conv1D) or isinstance(
        module, Depthwise_Separable_TransposeConv1D
    ):
        module.weight_norm()
        return module
    else:
        return weight_norm(module, name, dim)
//  [python]
def weight_norm_modules(module, name="weight", dim=0):
    if isinstance(module, Depthwise_Separable_Conv1D) or isinstance(
        module, Depthwise_Separable_TransposeConv1D
    ):
        module.weight_norm()
        return module
    else:
        return torch.nn.utils.parametrizations.weight_norm(module, name, dim)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, subbands, gin_channels=0):
        super(Multistream_iSTFT_Generator, self).__init__()
        # self.h = h
        self.subbands = subbands
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = weight_norm(
            Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u+1-i)//2, output_padding=1-i)))  # 这里k和u不是成倍数的关系，对最终结果很有可能是有影响的，会有checkerboard artifacts的现象

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.post_n_fft = gen_istft_n_fft
        self.ups.apply(commons.init_weights)
        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
        self.reshape_pixelshuffle = []

        self.subband_conv_post = weight_norm(
            Conv1d(ch, self.subbands*(self.post_n_fft + 2), 7, 1, padding=3))

        self.subband_conv_post.apply(commons.init_weights)

        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size

        updown_filter = torch.zeros(
            (self.subbands, self.subbands, self.subbands)).float()
        for k in range(self.subbands):
            updown_filter[k, k, 0] = 1.0
        self.register_buffer("updown_filter", updown_filter)
        self.multistream_conv_post = weight_norm(
            Conv1d(4, 1, kernel_size=63, bias=False, padding=commons.get_padding(63, 1)))
        self.multistream_conv_post.apply(commons.init_weights)
        self.cond = nn.Conv1d(256, 512, 1)
//  [python]
def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gen_istft_n_fft, gen_istft_hop_size, subbands, gin_channels=0):
        super(Multistream_iSTFT_Generator, self).__init__()
        # self.h = h
        self.subbands = subbands
        self.num_kernels = len(resblock_kernel_sizes)
        self.num_upsamples = len(upsample_rates)
        self.conv_pre = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3))
        resblock = modules.ResBlock1 if resblock == '1' else modules.ResBlock2

        self.ups = nn.ModuleList()
        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
            self.ups.append(torch.nn.utils.parametrizations.weight_norm(
                ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)),
                                k, u, padding=(k-u+1-i)//2, output_padding=1-i)))  # 这里k和u不是成倍数的关系，对最终结果很有可能是有影响的，会有checkerboard artifacts的现象

        self.resblocks = nn.ModuleList()
        for i in range(len(self.ups)):
            ch = upsample_initial_channel//(2**(i+1))
            for j, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
                self.resblocks.append(resblock(ch, k, d))

        self.post_n_fft = gen_istft_n_fft
        self.ups.apply(commons.init_weights)
        self.reflection_pad = torch.nn.ReflectionPad1d((1, 0))
        self.reshape_pixelshuffle = []

        self.subband_conv_post = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(ch, self.subbands*(self.post_n_fft + 2), 7, 1, padding=3))

        self.subband_conv_post.apply(commons.init_weights)

        self.gen_istft_n_fft = gen_istft_n_fft
        self.gen_istft_hop_size = gen_istft_hop_size

        updown_filter = torch.zeros(
            (self.subbands, self.subbands, self.subbands)).float()
        for k in range(self.subbands):
            updown_filter[k, k, 0] = 1.0
        self.register_buffer("updown_filter", updown_filter)
        self.multistream_conv_post = torch.nn.utils.parametrizations.weight_norm(
            Conv1d(4, 1, kernel_size=63, bias=False, padding=commons.get_padding(63, 1)))
        self.multistream_conv_post.apply(commons.init_weights)
        self.cond = nn.Conv1d(256, 512, 1)
//  [python]
def __init__(self, parameter, condition_dim=1024, hidden_dim=128, max_scale=1):
        super().__init__()
        self.parameter_shape = parameter.shape

        if len(self.parameter_shape) == 2:
            self.conditioners = torch.nn.Sequential(
                torch.nn.utils.weight_norm(torch.nn.Linear(condition_dim, hidden_dim)),
                torch.nn.Tanh(),
                torch.nn.utils.weight_norm(
                    torch.nn.Linear(
                        hidden_dim, 2 * (parameter.shape[0] + parameter.shape[1]) + 1
                    )
                ),
            )
        elif len(self.parameter_shape) == 1:
            self.conditioners = torch.nn.Sequential(
                torch.nn.utils.weight_norm(torch.nn.Linear(condition_dim, hidden_dim)),
                torch.nn.Tanh(),
                torch.nn.utils.weight_norm(
                    torch.nn.Linear(hidden_dim, 2 * parameter.shape[0] + 1)
                ),
            )
        else:
            raise RuntimeError()

        self.max_scale = max_scale
//  [python]
def __init__(self, parameter, condition_dim=1024, hidden_dim=128, max_scale=1):
        super().__init__()
        self.parameter_shape = parameter.shape

        if len(self.parameter_shape) == 2:
            self.conditioners = torch.nn.Sequential(
                torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(condition_dim, hidden_dim)),
                torch.nn.Tanh(),
                torch.nn.utils.parametrizations.weight_norm(
                    torch.nn.Linear(
                        hidden_dim, 2 * (parameter.shape[0] + parameter.shape[1]) + 1
                    )
                ),
            )
        elif len(self.parameter_shape) == 1:
            self.conditioners = torch.nn.Sequential(
                torch.nn.utils.parametrizations.weight_norm(torch.nn.Linear(condition_dim, hidden_dim)),
                torch.nn.Tanh(),
                torch.nn.utils.parametrizations.weight_norm(
                    torch.nn.Linear(hidden_dim, 2 * parameter.shape[0] + 1)
                ),
            )
        else:
            raise RuntimeError()

        self.max_scale = max_scale
//  [python]
def __init__(self, upscale=5, group=4, *args, **kwargs):
        super(OverNet, self).__init__()
        wn = lambda x: torch.nn.utils.weight_norm(x)
        self.upscale = upscale

        # self.sub_mean = MeanShift((0.4488, 0.4371, 0.4040), sub=True)
        # self.add_mean = MeanShift((0.4488, 0.4371, 0.4040), sub=False)

        self.entry_1 = wn(nn.Conv2d(3, 64, 3, 1, 1))

        self.GDG1 = LDGs(64, 64, wn=wn)
        self.GDG2 = LDGs(64, 64, wn=wn)
        self.GDG3 = LDGs(64, 64, wn=wn)

        self.reduction1 = BasicConv2d(wn, 64*2, 64, 1, 1, 0)
        self.reduction2 = BasicConv2d(wn, 64*3, 64, 1, 1, 0)
        self.reduction3 = BasicConv2d(wn, 64*4, 64, 1, 1, 0)

        self.reduction = BasicConv2d(wn, 64*3, 64, 1, 1, 0)

        self.Global_skip = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Conv2d(64, 64, 1, 1, 0), nn.ReLU(inplace=True))

        self.upsample = UpsampleBlock(64, upscale=upscale,  wn=wn, group=group)

        self.exit1 = wn(nn.Conv2d(64, 3, 3, 1, 1))

        self.res_scale = Scale(1)
        self.x_scale = Scale(1)
//  [python]
def __init__(self, upscale=5, group=4, *args, **kwargs):
        super(OverNet, self).__init__()
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        self.upscale = upscale

        # self.sub_mean = MeanShift((0.4488, 0.4371, 0.4040), sub=True)
        # self.add_mean = MeanShift((0.4488, 0.4371, 0.4040), sub=False)

        self.entry_1 = wn(nn.Conv2d(3, 64, 3, 1, 1))

        self.GDG1 = LDGs(64, 64, wn=wn)
        self.GDG2 = LDGs(64, 64, wn=wn)
        self.GDG3 = LDGs(64, 64, wn=wn)

        self.reduction1 = BasicConv2d(wn, 64*2, 64, 1, 1, 0)
        self.reduction2 = BasicConv2d(wn, 64*3, 64, 1, 1, 0)
        self.reduction3 = BasicConv2d(wn, 64*4, 64, 1, 1, 0)

        self.reduction = BasicConv2d(wn, 64*3, 64, 1, 1, 0)

        self.Global_skip = nn.Sequential(nn.AdaptiveAvgPool2d(1), nn.Conv2d(64, 64, 1, 1, 0), nn.ReLU(inplace=True))

        self.upsample = UpsampleBlock(64, upscale=upscale,  wn=wn, group=group)

        self.exit1 = wn(nn.Conv2d(64, 3, 3, 1, 1))

        self.res_scale = Scale(1)
        self.x_scale = Scale(1)
//  [python]
def reset_parameters(self) -> None:
        model_dim, kernel_size = self.in_channels, self.kernel_size[0]

        try:
            weight = self.weight_g

            remove_weight_norm(self)
        # Raised during the `__init__` call since we don't have the weight norm
        # hook registered yet. Safe to ignore.
        except AttributeError:
            weight = self.weight

            if weight.dtype == torch.bfloat16 and not torch_greater_or_equal(2, 2):
                raise RuntimeError(
                    "`torch.nn.utils.weight_norm()` supports `torch.bfloat16` only in PyTorch 2.2 and later versions."
                )

        nn.init.normal_(
            self.weight, mean=0.0, std=(4.0 / (kernel_size * model_dim)) ** 0.5
        )

        with catch_warnings():
            warnings.simplefilter("ignore")  # Suppress the deprecation warning.

            weight_norm(self, dim=2)

        self.weight_v.requires_grad_(weight.requires_grad)
        self.weight_g.requires_grad_(weight.requires_grad)

        if self.bias is not None:
            nn.init.constant_(self.bias, 0.0)
//  [python]
def reset_parameters(self) -> None:
        model_dim, kernel_size = self.in_channels, self.kernel_size[0]

        try:
            weight = self.weight_g

            remove_weight_norm(self)
        # Raised during the `__init__` call since we don't have the weight norm
        # hook registered yet. Safe to ignore.
        except AttributeError:
            weight = self.weight

            if weight.dtype == torch.bfloat16 and not torch_greater_or_equal(2, 2):
                raise RuntimeError(
                    "`torch.nn.utils.weight_norm()` supports `torch.bfloat16` only in PyTorch 2.2 and later versions."
                )

        nn.init.normal_(
            self.weight, mean=0.0, std=(4.0 / (kernel_size * model_dim)) ** 0.5
        )

        with catch_warnings():
            warnings.simplefilter("ignore")  # Suppress the deprecation warning.

            parametrizations.weight_norm(self, dim=2)

        self.weight_v.requires_grad_(weight.requires_grad)
        self.weight_g.requires_grad_(weight.requires_grad)

        if self.bias is not None:
            nn.init.constant_(self.bias, 0.0)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
        super(WN, self).__init__()
        assert(kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
        super(WN, self).__init__()
        assert(kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2*hidden_channels*n_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2*hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), rank=0):
        super(AMPBlock, self).__init__()

        self.convs1 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)
        self.alpha1 = nn.Parameter(torch.ones(1, channels, 1).to(rank))
        self.alpha2 = nn.Parameter(torch.ones(1, channels, 1).to(rank))

        self.convs2 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), rank=0):
        super(AMPBlock, self).__init__()

        self.convs1 = nn.ModuleList([
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)
        self.alpha1 = nn.Parameter(torch.ones(1, channels, 1).to(rank))
        self.alpha2 = nn.Parameter(torch.ones(1, channels, 1).to(rank))

        self.convs2 = nn.ModuleList([
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
        super(WN, self).__init__()
        assert (kernel_size % 2 == 1)
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size,
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout

        self.in_layers = torch.nn.ModuleList()
        self.res_skip_layers = torch.nn.ModuleList()
        self.drop = nn.Dropout(p_dropout)

        if gin_channels != 0:
            cond_layer = torch.nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
            self.cond_layer = torch.nn.utils.parametrizations.weight_norm(cond_layer, name='weight')

        for i in range(n_layers):
            dilation = dilation_rate ** i
            padding = int((kernel_size * dilation - dilation) / 2)
            in_layer = torch.nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size,
                                       dilation=dilation, padding=padding)
            in_layer = torch.nn.utils.parametrizations.weight_norm(in_layer, name='weight')
            self.in_layers.append(in_layer)

            # last one is not necessary
            if i < n_layers - 1:
                res_skip_channels = 2 * hidden_channels
            else:
                res_skip_channels = hidden_channels

            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
            res_skip_layer = torch.nn.utils.parametrizations.weight_norm(res_skip_layer, name='weight')
            self.res_skip_layers.append(res_skip_layer)
//  [python]
def __init__(self, args):
        super(EDRN, self).__init__()
        r = args.scale[0]
        G0 = args.G0
        kSize = args.EDRNkSize
        rgb_mean = (0.4313, 0.4162, 0.3861)
        rgb_std = (1.0, 1.0, 1.0)
        self.sub_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std)
        wn = lambda x: torch.nn.utils.weight_norm(x)
        # number of RG blocks, conv layers, out channels
        self.D, C, G = {
            'B': (4, 10, 16),
        }[args.EDRNconfig]

        self.SFENet1 = wn(nn.Conv2d(args.n_colors, G0, kSize, padding=(kSize-1)//2, stride=1))
        self.encoder1=nn.Sequential(*[
        wn(nn.Conv2d(G0, 2*G0, kSize, padding=(kSize-1)//2, stride=2)),
        nn.BatchNorm2d(2*G0),
        nn.ReLU(inplace=True)])
        self.encoder2=nn.Sequential(*[
        wn(nn.Conv2d(2*G0, 4*G0, kSize, padding=(kSize-1)//2, stride=2)),
        nn.BatchNorm2d(4*G0),
        nn.ReLU(inplace=True)
        ])
        self.decoder1=nn.Sequential(*[
        wn(nn.ConvTranspose2d(4*G0, 2*G0, 3, padding=1, output_padding=1, stride=2)),
        nn.BatchNorm2d(2*G0),
        nn.ReLU(inplace=True)])
        self.decoder2=nn.Sequential(*[
        wn(nn.ConvTranspose2d(2*G0, G0, 3, padding=1, output_padding=1, stride=2)),
        nn.BatchNorm2d(G0),
        nn.ReLU()
        ])
        RGs0 = [
            RG(growRate0 = 4*G0, nConvLayers = C, wn=wn) \
            for _ in range(self.D)]
        RGs0.append(wn(nn.Conv2d(4*G0, 4*G0, kSize, padding=(kSize-1)//2, stride=1)))
        RGs1 = [
            RG(growRate0 = 2*G0, nConvLayers = C, wn=wn) \
            for _ in range(self.D//2)]
        RGs1.append(wn(nn.Conv2d(2*G0, 2*G0, kSize, padding=(kSize-1)//2, stride=1)))
        RGs2 = [
            RG(growRate0 = G0,nConvLayers = C, wn=wn) \
            for _ in range(self.D//4)]
        RGs2.append(wn(nn.Conv2d(G0, G0, kSize, padding=(kSize-1)//2, stride=1)))
        self.RGs0 = nn.Sequential(*RGs0)
        self.RGs1 = nn.Sequential(*RGs1)
        self.RGs2 = nn.Sequential(*RGs2)
        self.restoration=wn(nn.Conv2d(G0, args.n_colors, kSize, padding=(kSize-1)//2, stride=1))
        self.add_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
//  [python]
def __init__(self, args):
        super(EDRN, self).__init__()
        r = args.scale[0]
        G0 = args.G0
        kSize = args.EDRNkSize
        rgb_mean = (0.4313, 0.4162, 0.3861)
        rgb_std = (1.0, 1.0, 1.0)
        self.sub_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std)
        wn = lambda x: torch.nn.utils.parametrizations.weight_norm(x)
        # number of RG blocks, conv layers, out channels
        self.D, C, G = {
            'B': (4, 10, 16),
        }[args.EDRNconfig]

        self.SFENet1 = wn(nn.Conv2d(args.n_colors, G0, kSize, padding=(kSize-1)//2, stride=1))
        self.encoder1=nn.Sequential(*[
        wn(nn.Conv2d(G0, 2*G0, kSize, padding=(kSize-1)//2, stride=2)),
        nn.BatchNorm2d(2*G0),
        nn.ReLU(inplace=True)])
        self.encoder2=nn.Sequential(*[
        wn(nn.Conv2d(2*G0, 4*G0, kSize, padding=(kSize-1)//2, stride=2)),
        nn.BatchNorm2d(4*G0),
        nn.ReLU(inplace=True)
        ])
        self.decoder1=nn.Sequential(*[
        wn(nn.ConvTranspose2d(4*G0, 2*G0, 3, padding=1, output_padding=1, stride=2)),
        nn.BatchNorm2d(2*G0),
        nn.ReLU(inplace=True)])
        self.decoder2=nn.Sequential(*[
        wn(nn.ConvTranspose2d(2*G0, G0, 3, padding=1, output_padding=1, stride=2)),
        nn.BatchNorm2d(G0),
        nn.ReLU()
        ])
        RGs0 = [
            RG(growRate0 = 4*G0, nConvLayers = C, wn=wn) \
            for _ in range(self.D)]
        RGs0.append(wn(nn.Conv2d(4*G0, 4*G0, kSize, padding=(kSize-1)//2, stride=1)))
        RGs1 = [
            RG(growRate0 = 2*G0, nConvLayers = C, wn=wn) \
            for _ in range(self.D//2)]
        RGs1.append(wn(nn.Conv2d(2*G0, 2*G0, kSize, padding=(kSize-1)//2, stride=1)))
        RGs2 = [
            RG(growRate0 = G0,nConvLayers = C, wn=wn) \
            for _ in range(self.D//4)]
        RGs2.append(wn(nn.Conv2d(G0, G0, kSize, padding=(kSize-1)//2, stride=1)))
        self.RGs0 = nn.Sequential(*RGs0)
        self.RGs1 = nn.Sequential(*RGs1)
        self.RGs2 = nn.Sequential(*RGs2)
        self.restoration=wn(nn.Conv2d(G0, args.n_colors, kSize, padding=(kSize-1)//2, stride=1))
        self.add_mean = MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
//  [python]
def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1):
        super().__init__()
        self.padding = dilation * (kernel_size - 1) // 2
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size,
                              dilation=dilation, padding=self.padding)
        self.conv = nn.utils.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)
//  [python]
def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1):
        super().__init__()
        self.padding = dilation * (kernel_size - 1) // 2
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size,
                              dilation=dilation, padding=self.padding)
        self.conv = nn.utils.parametrizations.weight_norm(self.conv)
        nn.init.kaiming_normal_(self.conv.weight)
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), rank=0, orig_freq=None, rolloff=0.25):
        super(AMPBlock_filter, self).__init__()

        self.upsampling_with_lfilter = T.Resample(orig_freq=orig_freq, new_freq=orig_freq * 2,
                                                  resampling_method='kaiser_window',
                                                  lowpass_filter_width=12,
                                                  rolloff=rolloff,
                                                  beta=4.663800127934911
                                                  )
        self.downsampling_with_lfilter = T.Resample(orig_freq=orig_freq * 2, new_freq=orig_freq,
                                                    resampling_method='kaiser_window',
                                                    lowpass_filter_width=12,
                                                    rolloff=rolloff,
                                                    beta=4.663800127934911
                                                    )
        self.convs1 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList([
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)

        self.alpha1 = [nn.Parameter(torch.ones(1, channels, 1).to(rank)) for i in range(len(self.convs1))]
        self.alpha2 = [nn.Parameter(torch.ones(1, channels, 1).to(rank)) for i in range(len(self.convs2))]
//  [python]
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5), rank=0, orig_freq=None, rolloff=0.25):
        super(AMPBlock_filter, self).__init__()

        self.upsampling_with_lfilter = T.Resample(orig_freq=orig_freq, new_freq=orig_freq * 2,
                                                  resampling_method='kaiser_window',
                                                  lowpass_filter_width=12,
                                                  rolloff=rolloff,
                                                  beta=4.663800127934911
                                                  )
        self.downsampling_with_lfilter = T.Resample(orig_freq=orig_freq * 2, new_freq=orig_freq,
                                                    resampling_method='kaiser_window',
                                                    lowpass_filter_width=12,
                                                    rolloff=rolloff,
                                                    beta=4.663800127934911
                                                    )
        self.convs1 = nn.ModuleList([
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
                               padding=get_padding(kernel_size, dilation[0]))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
                               padding=get_padding(kernel_size, dilation[1]))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
                               padding=get_padding(kernel_size, dilation[2])))
        ])
        self.convs1.apply(init_weights)

        self.convs2 = nn.ModuleList([
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1))),
            parametrizations.weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
                               padding=get_padding(kernel_size, 1)))
        ])
        self.convs2.apply(init_weights)

        self.alpha1 = [nn.Parameter(torch.ones(1, channels, 1).to(rank)) for i in range(len(self.convs1))]
        self.alpha2 = [nn.Parameter(torch.ones(1, channels, 1).to(rank)) for i in range(len(self.convs2))]
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
        sigmoid_scale=False,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        self.sigmoid_scale = sigmoid_scale

        start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
        start = torch.nn.utils.weight_norm(start)
        self.start = start
        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  It helps to stabilze training.
        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        self.wn = modules.WN(
            in_channels,
            hidden_channels,
            kernel_size,
            dilation_rate,
            n_layers,
            gin_channels,
            p_dropout,
        )
//  [python]
def __init__(
        self,
        in_channels,
        hidden_channels,
        kernel_size,
        dilation_rate,
        n_layers,
        gin_channels=0,
        p_dropout=0,
        sigmoid_scale=False,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        self.sigmoid_scale = sigmoid_scale

        start = torch.nn.Conv1d(in_channels // 2, hidden_channels, 1)
        start = torch.nn.utils.parametrizations.weight_norm(start)
        self.start = start
        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  It helps to stabilze training.
        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        self.wn = modules.WN(
            in_channels,
            hidden_channels,
            kernel_size,
            dilation_rate,
            n_layers,
            gin_channels,
            p_dropout,
        )
//  [python]
def conv_bn(inp, oup, stride):
    return nn.Sequential(
        torch.nn.utils.weight_norm(nn.Conv2d(inp, oup, 3, stride, 1, bias=False)),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
    )
//  [python]
def conv_bn(inp, oup, stride):
    return nn.Sequential(
        torch.nn.utils.parametrizations.weight_norm(nn.Conv2d(inp, oup, 3, stride, 1, bias=False)),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True)
    )
//  [python]
def __init__(self, in_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, sigmoid_scale=False, n_sqz=1, shared_wn=None):
        super().__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels
        self.kernel_size = kernel_size
        self.dilation_rate = dilation_rate
        self.n_layers = n_layers
        self.gin_channels = gin_channels
        self.p_dropout = p_dropout
        self.sigmoid_scale = sigmoid_scale

        start = torch.nn.Conv1d(in_channels//2, hidden_channels, 1)
        start = torch.nn.utils.weight_norm(start)
        self.start = start
        # Initializing last layer to 0 makes the affine coupling layers
        # do nothing at first.  It helps to stabilze training.
        end = torch.nn.Conv1d(hidden_channels, in_channels, 1)
        end.weight.data.zero_()
        end.bias.data.zero_()
        self.end = end

        self.cond_layer = torch.nn.utils.weight_norm(
            torch.nn.Conv1d(
                n_sqz * gin_channels, 2*hidden_channels*n_layers, 1), name='weight')

        # Grouped Parameter Sharing
        self.wn = shared_wn if shared_wn is not None else NonCausalWaveNet(hidden_channels, kernel_size,
                                                                           dilation_rate, n_layers, gin_channels, p_dropout, n_sqz)