I ended up setting up that linear layer to contain a vector of weights rather than a matrix. Then when it comes time to apply the vector I simply do torch.diag_embed(vector) (which is differentiable) to get the diagonal matrix. This way the weight matrix is only ever diagonal.