I am building a model that outputs a scaler, and I need to compute the gradient of those outputs with respect to their inputs. I am doing something like the following:


def get_forces(outputs, coordinates):
    forces = []
    for index, (hash, image) in enumerate(coordinates.items()):
        for i, coordinate in image: 
            forces_ = -torch.autograd.grad(
                                      outputs[index], 
                                      coordinate,  
                                      create_graph=True,
                                      retain_graph=True
             )[0]
            forces.append(forces_)
    return torch.stack(forces)


while not converged:
    optimizer.zero_grad()  # clear previous gradients
    outputs = model(inputs)    # Note that inputs have requires_grad = True
    outputs = torch.stack(list(outputs.values())).sum(0)
    

    all_grad = get_gradients(outputs, coordinates[0])

    outputs = {"scalar": outputs, "gradiend": all_grad}

    criterion = CustomLossFunction(gradient=True)
    loss = criterion(outputs, targets)    # This loss uses both scalar and gradient
    
    loss.backward()  
    optimizer.step()

    epoch += 1

    
    if epoch == stop:
        converged = True

The first epoch passes correctly, but at the second epoch I get:
RuntimeError: Trying to backward through the graph a second time, but 
he saved intermediate results have already been freed. Specify 
retain_graph=True when calling backward the first time.

How can this be solved? How can I compute the gradient of the output of my neural network with respect to its inputs without having this issue? Meanwhile, I am reading on the forum about how other people have solved this.