Assuming loss1 and loss2 are supposed to create the gradients for modelA and modelB, you could use loss3.backward(retain_graph=True, inputs=modelB.parameters()) and call loss.backward() on the accumulated loss from loss1 and loss2 afterwards.



 hsc:

I could do it using retain_graph=True, but it requires more gpu memory.


It would require more memory if you don’t free the computation graph in the last backward call, which shouldn’t be the case.