andrea.migliorati:

and the torch.autograd.grad function, which seems to work on scalars only instead


Hi Andrea!
torch.autograd.grad isn’t restricted to only scalar outputs, you just need to modify the grad_outputs argument of torch.autograd.grad. So, let’s say you have a function which is N-dimensional input and an M-dimensional output and a batch size of 100. You just need to add torch.ones_like(Y) as your 3rd argument to torch.autograd.grad
B=100
N=4
M=6

X = torch.randn(B, N, requires_grad=True)
fc = torch.nn.Linear(N,M) #our R^N -> R^M function 

Y = fc(X)

dYdX = torch.autograd.grad(Y, X, torch.ones_like(Y)) #gradient of the function
#returns torch.Size([100, 4])