Dear Thomas
Thanks for your posting. I am also encountering the ‘optimize non-leaf variable’ issue and I am grategul if you can provide some feedback to my practice.
What I want to do is using a linear combination (learnable) of  basis filters for convolution operation and optimize only the linear coefficients as below:
class net(torch.nn.Module):
    def __init__():
        fitler_basis = [f_height, f_weight, num_filter];
        self.coeff = [num_filter, c_in, c_out]; # To be optimized

        self.conv2d = torch.nn.conv2d(c_in, c_out, strides=1, padding=1)
        self.conv2d.weight = torch.matmul(filter_basis, coeff) # linear combinaton of pre-fixed filter as weights;

    def forward(self, data):
        return self.conv2d(data)

    def train(self):
        loss = ...
        optim(self.coeff, loss)

How should I achieve my goal ?