In the end, you mathematically only get a gradient of the shape of your parameters if you differentiate a scalar function.
There are cases where you want the Jacobian, ie differentiate a vector output, but those are rare enough that PyTorch does not support it as well as differentiating scalars. (But you’ll find tricks if you search for Jacobian here on the forum.)
Best regards
Thomas