Right now, I’m doing:
output_gradients = []

for output in net_outputs:
    tmp_grad = {}
    net.zero_grad()
    output.backward(retain_graph=True)
    for name, param in net.named_parameters():
         tmp_grad[name] = param.grad
    output_gradients.append(tmp_grad)

Since I have to call backward on each output, the backward passes are not parallelized and thus, the code is pretty slow.
Is there a faster way? Thanks!