Let’s say I have a cascaded model:
e.g., x → modelA → modelB → output <-loss-> gt
now, there are multiple losses (loss1, loss2, loss3) I’d like to calculate but I want to backpropagate only loss1 and loss2 to modelA, not loss3.
The modelB needs to be updated using every loss though.
I could do it using retain_graph=True, but it requires more gpu memory.
I feel like I can do this by masking the gradient of a certain loss function but I’m not sure how I can do this.
Thank you.