Based on your figure you would pass the output of F (let’s call it x_f) to the modules A and B.
To keep the computation graph for the A path you would simply pass it to A, calculate the loss, and call backward on it. However, since you only want to calculate the gradients for the parameters of B using this aux. path, you could detach x_f and pass it to B:
x_f = F(x)
x_a = A(x_f)
x_b = B(x_f.detach())
return x_a, x_b