Partial Backwards in PyTorch Graph
I have a medium sized tensor x. On this mediums sized tensor an computationally expensive function (both forward and backwards) q is applied to obtain another medium sized tensor y.
Using y I evaluate many functions f producing a scalar, they are not particularly computationally expensive however use large internal states causing a large computational graph for torch that.
Now I want to compute the gradient on x in the following way
y = q(x)

for f in functions
    res += f(y)

res.backward()

The proble with this implementation is that the graphs of all the functions f are retained. That causes the total memory usage to explode.
The other possibility would be to compute
y = q(x)

for f in functions
    partial = f(y)
    partial.backward(retain_graph = True)

The advantage is that after every function evaluation f the result goes out of scope and the graph is freed, saving massive amounts of memory. However, in this situation the function q(x) is evaluated backwards multiple times which is very time expensive.
In the ideal situation I would want to first compute the gradient for y using code similar to the second example and then only once evaluate q backwards to get the gradient for x. What is the proper way to do that with PyTorch?

I think that would be the way to achieve it :
y = q(x)
z = y.detach()
z.requires_grad_(True)

for f in functions:
    partial = f(y)
    partial.backward(retain_graph = True)
y.backward(z.grad)

You accumulate all the gradients in z, which y but in another computational graph, then you propagate these gradients (z.grad) in the first graph
