Weight Normalisation causes Non-Leaf error with ADAM optimiser Pytorch

I am using a PyTorch model (model given in the end).

When I use

optimizer = torch.optim.Adam(model.parameters(), lr=opt['base_lr'])

The code works fine.

However when I use,

def get_parameters(model, bias=False):
    for m in model.modules():
        if bias:
            if isinstance(m, nn.Conv2d):
              if m.bias is not None:
                yield m.bias
        else:
            if isinstance(m, nn.Conv2d):
              if m.weight is not None:
                yield m.weight


optimizer = torch.optim.Adam(
    [
        {'params': get_parameters(model, bias=False)},
        {'params': get_parameters(model, bias=True), 'lr': opt['base_lr'] * 2, 'weight_decay': 0},
    ],lr=opt['base_lr'])

A runtime error is thrown stating, ValueError: can't optimize a non-leaf Tensor

I understand the difference between leaf and non-leaf tensor and found that if I switch off weight normalisation in my model, no error is thrown even with the modified ADAM optimiser. i.e replacing layers.append(wn(m)) with layers.append(m) where wn is given by from torch.nn.utils import weight_norm as wn in my model definition prevents any error.

Indeed I checked using is_leaf and weight normalisation is making tensors Non-leaf.

Can someone help me how I can fix the problem while retaining Weight Normalisation?

My model is given below:

import torch
from torch.nn.utils import weight_norm as wn
import torch.nn as nn

def conv_layer(inc, outc, kernel_size=3, groups=1, bias=False, negative_slope=0.2, bn=False, init_type='kaiming', fan_type='fan_in', activation=False, pixelshuffle_init=False, upscale=4, num_classes=3, weight_normalization = True):

    layers = []
    
    if bn:
        m = nn.BatchNorm2d(inc)
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)
        layers.append(m)
    
    if activation=='before':
            layers.append(nn.LeakyReLU(negative_slope=negative_slope, inplace=False))
        
    if pixelshuffle_init:
        m = nn.Conv2d(in_channels=inc, out_channels=num_classes * (upscale ** 2),
                                  kernel_size=3, padding = 3//2, groups=1, bias=True, stride=1)
        nn.init.constant_(m.bias, 0)
        with torch.no_grad():
            kernel = ICNR(m.weight, upscale, negative_slope, fan_type)
            m.weight.copy_(kernel)
    else:
        m = nn.Conv2d(in_channels=inc, out_channels=outc,
     kernel_size=kernel_size, padding = (kernel_size-1)//2, groups=groups, bias=bias, stride=1)
        init_gain = 0.02
        if init_type == 'normal':
            torch.nn.init.normal_(m.weight, 0.0, init_gain)
        elif init_type == 'xavier':
            torch.nn.init.xavier_normal_(m.weight, gain = init_gain)
        elif init_type == 'kaiming':
            torch.nn.init.kaiming_normal_(m.weight, a=negative_slope, mode=fan_type, nonlinearity='leaky_relu')
        elif init_type == 'orthogonal':
            torch.nn.init.orthogonal_(m.weight)
        else:
            raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
    
    if weight_normalization: 
        layers.append(wn(m)) # <------ REMOVING THIS MAKES CODE EXECUTE WITH NO ERROR
    else:
        layers.append(m)
    
    if activation=='after':
            layers.append(nn.LeakyReLU(negative_slope=negative_slope, inplace=False))
            
    return nn.Sequential(*layers)

class Net(nn.Module):
    
    def __init__(self):
        super(Net, self).__init__()
        
        
        self.up4 = nn.PixelShuffle(4)
        self.up2 = nn.PixelShuffle(2)
        
        self.conv32x = nn.Sequential(        
                        conv_layer(1024, 128, kernel_size=3, groups=128, bias=True, negative_slope=1, bn=False, init_type='kaiming', fan_type='fan_in', activation=False, pixelshuffle_init=False, upscale=False, num_classes=False, weight_normalization = True),
                        conv_layer(128, 64, kernel_size=3, groups=1, bias=True, negative_slope=1, bn=False, init_type='kaiming', fan_type='fan_in', activation=False, pixelshuffle_init=False, upscale=False, num_classes=False, weight_normalization = True)
                        )
...