To compute gradients (which implies that the output of the function is scalar),
The outputs you pass to autograd.grad should be either of the following:

A Variable wrapping a Tensor of size (1,) (something like a scalar)
An arbitrary tuple of Variables. In this case, you should specify a grad_output= for autograd.grad that has the same shape as the arbitrary tuple of Variables.