What you just proposed gives me this error:
RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.

I would have used the grad_outputs parameter to avoid the loop:
x = torch.tensor([[1.], [1.]], requires_grad=True) 
y = x**2  # shape (2, 1)
y_x = grad(y, x, grad_outputs=torch.ones_like(y)) 

"""
tensor([[2.], [2.]])
"""

y = x*torch.tensor([[3.], [2.]]) 
y_x = grad(y, x, grad_outputs=torch.ones_like(y)) 
"""
tensor([[3.], [2.]])
"""