The thing is I’ll freeze one fc layer, and then train the model with second fc layer. But in second phase, I’ll train the model with both fc layers (unfreeze) on different dataset, so how will i compute loss pertaining to two fc layers unlike the code which I provided.