Thanks for the help! The trick for optimizer is so helpful !
Yes; my plan would be (1) freeze the layer I don’t want, train for few epochs, then unfreeze the network and (2)train the whole again.
So (1) and (2) have two loss functions here.
But in this case, did you mean that I must have  retain_graph=True  ? Because so far I must have both losses to retain_graph. This leads to out of memory