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 instead returns the gradients as Variable tuples.
grad_sum = 0
for grad in grads:
    grad_sum += grad.pow(2).sum()
# grad_sum.backward() will accumulate the gradients into the .grad attributes
grad_sum.backward(retain_graph=True)
optimizer.step()

In my case, instead of inputs, I used layer activations, but should work similar for model inputs.
EDIT: you don’t need the retain_graph=True in grad_sum.backwards call if you don’t intend to do anything else with these gradients before the weight update.