Hello everyone, sorry in advance for the beginner question.
Let’s say I have a small model net, with d parameters, with dimension c as output (like a classifier for MNIST). Let’s also say that I feed this model with a minibatch of images (batchsize B). At this point I would have something like
outputs = net(images)

where the shape of output will be (B, c). I would like to take the gradients with respect to the parameters of the network, something like this:
outputs.backward()
feature_maps = [param.grad for param in iter(net.parameters())]

And having as output this feature_maps with shape (B, c, d). This is not doable for many reasons I guess, in particular pytorch complains that when I call .backward() I need to have a scalar tensor in front of it.
This problem is solvable, I just have to feed one image at a time in the model, and compute the gradient for each output separately, and then putting everything together again. Is there a smarter/faster way to achieve this?
I’m basically computing the feature maps for computing the neural tangent kernel later (even if I’m doing it for each output separately for this classifier), so I believe there should be a better usual way for doing this.