@JuanFMontesinos
Here is my whole train process.
Can you check last?
class Network(nn.Module):
 ...
class NetworkToFreeze(nn.Module):
 ...

for p in network.parameters():
    p.requires_grad = True
for p in network_to_freeze.parameters():
    p.requires_grad = True
...
for epoch in train_process:
       if epoch < 50:
           train all network code(Same code here*)

           total_loss = network_loss + network_to_freeze_loss
           total_loss.backward(retain_graph=True)

       else:
           train all network code(Same code here*)

           for p in network_to_freeze.parameters():
               p.requires_grad = False

           total_loss = network_loss
           total_loss.backward(retain_graph=True)