I am developing a model where, for some technical reasons, my minimization objective is actually the squared norm of the gradient of some loss function:

In the following toy example, I am implementing this for a very simple case (least-squares linear regression from the 2-dimensional space to the 1-dimensional space):
import torch
from torch import Tensor
from torch.autograd import Variable
from torch import nn

# some toy data
x = Variable(Tensor([4., 2.]), requires_grad=False)
y = Variable(Tensor([1.]), requires_grad=False)

# linear model and squared difference loss
model = nn.Linear(2, 1)
loss = torch.sum((y - model(x))**2)

# first backward pass
model.zero_grad()
# set create_graph=True to allow the computation of higher order derivatives afterwards
loss.backward(create_graph=True)

# compute the squared norm of the loss gradient
gn2 = Variable(Tensor([0.]))
for param in model.parameters():
  gn2 = gn2 + param.grad.norm()**2

# we do not want to accumulate the previous gradients with the new ones,
# so we have to call zero_grad() again
model.zero_grad()
gn2.backward()


However, PyTorch complains about backpropagating through variables that were modified by inplace operations:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

After some debug, I found out that the second call to model.zero_grad() is responsible for that error, because it sets gradients to zero using inplace operations. An obvious (but maybe inelegant?) workaround is manually “zeroing” the gradients in the for loop and removing the second call to model.zero_grad(), that is, doing:
...
# compute the squared norm of the loss gradient
gn2 = Variable(Tensor([0.]))
for param in model.parameters():
  gn2 = gn2 + param.grad.norm()**2
  param.grad = torch.zeros_like(param)

gn2.backward()

Everything would be cleaner if the function  zero_grad() had an inplace argument, that we could set to False in situations like this…
What do you think? Which is the correct way to deal with this situation?
Thanks in advance.