Hi,
I’m afraid it’s not possible.
If you don’t use regularization, one way would be to set the gradients for the diagonal elements to 0.
But if you have regularization they will still change.
Otherwise, you can save your parameters diagonal manually before training and then re-modify the weights after your optimizer step to the original value.