Precomputing part of a loss function

I’m trying to use pytorch to calculate a complex function but still require the gradient of the output as a function of the inputs. For example:

a=torch.tensor([1,2,3], dtype=torch.float32, requires_grad=True)
b=torch.tensor([3,2,1], dtype=torch.float32, requires_grad=True)
v=torch.tensor([0.2], dtype=torch.float32, requires_grad=True)

# here precalc represents some (fairly expensive) sequence of operations
precalc =*b+a*a+b*b

def calc(precalc, v):
    return torch.relu(batch).mean()

So the precalc tensor is a fixed function of a and b (but expensive to compute).

When I call calc for the first time as x=calc(precalc, v) and then x.backward(), I get the correct gradients for a,b and v. However, if I call calc a second time, i.e. x2=calc(precalc, v) and then x2.backward(), I get the pytorch error that I’m going backward through the graph a second time and I should retain the graph.

I would ideally like to free the graph as precalc is expensive to compute but small in memory and the calc itself is easier to compute but will use a lot of memory.

Is there a way to achieve this?