Does It make sense ?
   if epoch < 50:
       train all network code(Same code here*)
       total_loss = network_loss + network_to_freeze_loss
       total_loss.backward()
   if epoch >=50:
       for p in network_to_freeze.parameters():
           p.requires_grad = False
       train all network code(Same code here*)
       total_loss = network_loss
       total_loss.backward()