I scrapped working on this back when I posted it, and only restarted a week ago. I found your ReparamModule really helpful. However, I would like to understand one thing. Using your class def I wrote a simple module as follows.
@add_metaclass(PatchModules)
class ReparamModule(nn.Module):
    def __init__(self):
        super(ReparamModule, self).__init__()
        self.fc1 = nn.Linear(10, 1)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        return x
   #other functions 

rpm = ReparamModule()

Also, I did not assign the old weights back as you did in



 SimonW:

https://gist.github.com/SsnL/f2f56534aefb22d8612dbd7a5da28ed8#file-reparam_module-py-L51


since I want the gradients to be computed with respect to the newly assigned parameters new_ws.
I find that after forward_with_weights, rpm.fc1.weight, rpm.fc1.bias have the assigned weights, but if I call rpm.get_weights() I get the old weights. I see that the parameters are also unchanged. Similarly, I had to explicitly use rpm..fc1.weight.grad instead of looping over rpm.parameters() and use .grad.data. Why is this so? How are the parameters and weight, bias decoupled? I attached the code for your reference.
https://gist.github.com/JayanthRR/07aa7a8d45a027954a8d0d2bb1fbaec7 3