Hi,
I have net_A (+) which I want to train, and net_B (-) which is pre-trained and has requires_grad=False.
net_C is as follows:
+++++++++++±----------
net_A  -> net_B (iterate N times)
Sample code:
for epoch in range(arg.epochs):
    for i, data in enumerate(dataloader):
        net_A.zero_grad()
        # data = batch x 3 x 28 x 28
        A = net_A(data)
        # A = batch x N x 2
        for b in range(A.size(0)):
            for n in range(A.size(1)):
                B = net_B(A[b, n:n+1])
                loss = A_loss(B, data) # some loss that involves net_B and input data
                loss.backward(retain_graph=True)
        optimizer.step() # net_A parameters

In the computation graph net_A (+) is trained and needed for every backprop, but net_B (-) branches (leafs) are changed.
retain_graph saves all (-) paths, which are not needed (only the current leaf is needed) and ends up returning “out of memory”…
How can I retain only the sub-graph (+) ?
Thanks