Can compute_grad1 support multiple backprops?

[zhangxz1123]

Hi Yaroslav,
First of all, I want to say your tricks here are really neat and have helped me a lot!

I'm currently trying to compute the per-example gradient over two backprops in a typical Hessian vector product calculation, e.g.

output = model(data)
loss = loss_fn(output, targets)
grad = torch.autograd.grad(loss, model.parameters(), create_graph=True)
grad = flat_grad(grad)
grad_p = (grad * p).sum()
hessian_p= torch.autograd.grad(grad_p , actor.parameters())

Is it possible to apply compute_grad1 to this procedure to get the per-example Hessian vector product? Naively applying it fails because currently compute_grad1 only supports a single backprop.

Thanks!
Jack