Hi, thanks for your answer! I’ll try to give an example of what I’m doing in pseudocode:
Let’s say I have 32 samples in a batch which I run through a module net1 that maps them to size 10:
[32 x 10]
net1_out = net1(samples)

Now I want the hessian of another module net2 w.r.t. these:
hessian(net2, net1_out, create_graph=True)
What I mean by “in a batch” is that I would like the output to have shape: [32 x 10 x 10]
However, the output seems to be: [32 x 10 x 32 x 10]
In other words, it seems to be computing the hessians across batch items, which is more expensive than what I need.