Hi,
Thanks for the response.
Just to make sure I was replacing weights and biases the right way, I commented the reduce(w,b) call and made the below changes to the code
for layer_idx in range(len(layers)):
    w = layers[layer_idx].weight.data.clone()
    b = layers[layer_idx].bias.data.clone()
    # new_w, new_b = reduce(w, b) # Custom logic to calculate new sets of w and b (always new_b.shape, w.shape)

    layers[layer_idx].weight = nn.Parameter(w)
    layers[layer_idx].weight.data = nn.Parameter(b)
    layers[layer_idx].bias = nn.Parameter(w)
    layers[layer_idx].bias.data = nn.Parameter(b)
    

    # Changing the grad values to match the shape of new weights and bias. 
    # Assigning random values to grad. This is just to match shapes 
	layers[layer_idx].weight.grad = nn.Parameter(torch.ones(w.shape))
    layers[layer_idx].weight.data.grad = nn.Parameter(torch.ones(w.shape))    
    layers[layer_idx].bias.grad = nn.Parameter(torch.ones(b.shape))
	layers[layer_idx].bias.data.grad = nn.Parameter(torch.ones(b.shape))

Basically I am not making any changes to weights and bias. I re-assign them to see if the network trains normally.
However, to my surprise, gradients after first epoch never change. All the gradient values equal to 1. By debugging optimizer’s param_groups I see that for every layer the gradients for weight and bias are 1. Although I initialize it to 1 after every epoch(which I do to match the shape), the loss.backward() and optimizer.step() should calculate new gradients right ?
Is there anything I might be missing? Please point me in the right direction if I am doing something stupid.
Thank You