To speed up calculating the gradient of output w.r.t. input, I am trying to use the minibatch. But batching seems to be a problem for autograd(), Is there any way to get the gradients?

Here are the details:

Input: (Batch, C1, C2, D)

Output: (Batch, ) Each item in the output is a probability value (scalar)

How to calculate the gradient of Output w.r.t Input? Such that the gradient should be of size (Batch, C1, C2, D).

Note that it can be easily achieved with the for loop, but I want it to be more efficient.

Thanks!