I’m trying to generate multiple saliency maps for a model with multiple outputs and a single input image. I’m calling torch.autograd.grad once for each of the outputs as I want seperate gradients for each output.
My problem is that when I call torch.autograd.grad multiple times, each call increases the GPU memory usage. Here’s a code snippit:
    input_gradients = []
    for output_scalar in multiple_output_scalars:
        self.model.zero_grad() #not needed right?
        input_gradients.append( torch.autograd.grad(
            outputs = output_scalar, 
            inputs = x, 
            retain_graph=True,
            create_graph=False)[0].cpu().detach() )
    input_gradients = torch.cat(input_gradients,dim=0)

I probably misunderstand what auto grad is doing, but to my understanding it should be using the computation graph to compute the gradiens and returning the gradients. Nothing new should be stored and since the graph is retained, nothing needs rebuilt. But clearly I’m wrong. Where is my thinking flawed and what is the right way to do this?