Hello,
please consider the following code:
There are multiple layers before this affine operation and activation layer
    self.fc1 = nn.Linear(13824, 4096) 
    torch.nn.init.kaiming_normal_(self.fc1.weight) 
    self.batchnorm_fc1 = nn.BatchNorm1d(4096)

    #YL
    self.fc3 = nn.Linear(4096,  3* 11* 13)
    torch.nn.init.kaiming_normal_(self.fc3.weight)
    self.batchnorm_fc3 = nn.BatchNorm1d( 3* 11* 13)
    
    #YH[0]
    self.fc4 = nn.Linear(4096,   3* 3* 29* 39)
    torch.nn.init.kaiming_normal_(self.fc4.weight)
    self.batchnorm_fc4 = nn.BatchNorm1d( 3* 3* 29* 39)
    
    #YH[1] there are two more layers here 
    #YH[2] 

I start creating multiple heads here as the output, however YH[0] for instance has a product of numbers (you might call it weights) around 10k while YL has product of numbers of around ~400. I would like to treat each of those layers (and their weight) differently when it comes to their learning rates, otherwise the gradient is propagating through all those 400+10000+… numbers which overlooks the significance of YL sub-band output which is the most significant in my case. What do you recommend?!
with the forward propagation:
    x = F.leaky_relu( self.batchnorm_fc1(self.fc1(x)) )
    ###############

           
    xl = self.fc3(x)
    xl  = xl .view(-1, 3, 11, 13)
    
    xh=[]
    xh.append(self.fc4(x))
    xh[0] = xh[0].view(-1,3, 3, 29, 39)
    ##same done with the xh[1] and xh[2]

      
    # return the output
    return xl,xh

Could you guide me on how to handle xl, and xh’s layers weights differently?!
I am using those 4 output heads for constructing an inverse wavelet transform(2d) but that isn’t so important right now.