Auto gradient for a batch of input

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.


1 Like


The autograd has no notion of “batch” and treats all dimensions the same.
So if you have a tensor that requires_gradient and is a leaf of size batch x C1 x C2 xD, then it will compute the gradients for it properly. You only need to adapt your ops in the middle to handle properly this new op. Most pytorch op should do that out of the box :slight_smile:

Thanks for the reply! That’s true~