I wonder how to set up an NN module whose parameters are non-leaf variables.
What I would like to do is something like below:
class NNTest(torch.nn.Module):
    def __init__(self):
        super(NNTest, self).__init__()
        self.module = torch.nn.Linear(2, 3, bias=False)
    
    def forward(self, x):
        return self.module(x)

nn = NNTest()
p = torch.nn.Parameter(torch.randn(2, 3))
d = p**2
nn.module.weight = d  # raises a TypeError here
# nn.module.weight.data = d.data  # this is not what I want

The problem of the last line is that it does not propagate the gradient to p and I would like the gradient to propagate beyond the Module’s weight.
Any help will be appreciated! Thanks!