Solved by michaelklachko in post #2 
              
              
                Here’s how I implemented a penalty on gradient growth, perhaps this will be useful to you: 
optimizer.zero_grad()
loss.backward(retain_graph=True)
grads = torch.autograd.grad(loss, inputs, create_graph=True)
# torch.autograd.grad does not accumuate the gradients into the .grad attributes 
# it inste…