Hi, I am trying to acquire the gradient in the last hidden layer for batch inputs.
Let says the last layer is a linear layer (512*100) and 100 is the number of classes.
For each sample, the gradient would be 100*512
Currently, my code is as follows.
losses = self.criterion_without_reduction(outputs, targets)
gradient_batch = []
for loss in losses:
    loss.backward(retain_graph=True)
    gradient_list = reversed(list(model.net.named_parameters()))
    # only output the gradient for the last layer
    for name, parameter in gradient_list:
        if 'weight' in name:
            gradient = parameter.grad.clone() # [column[:, None], row].resize_(100,100)
            gradient = gradient.unsqueeze_(0)
            gradient_batch.append(gradient.unsqueeze_(0))

            break
gradient_batch = torch.cat(gradient_batch, dim=0)

The above code will generate a tensor ([batch_size, 1, 100, 512])
However, the code runs relatively slow and I am wondering whether there is a more efficient way to do it?