Solved by tom in post #2 
              
              
                After loss.backward() and before optimizer.step(), you can do 
model.layer.weight.grad.diagonal().zero_(). 
Having gradient entries constant to zero on every invocation should leave the values fixed with most optimizers. 
Sophisticated people might mention backward hooks, but I am not one of them. 
B…