Hi,
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