Freezing partial weights after n epochs

[johnathanchiu]

I want to be able to freeze a portion of my model after n epochs of training. I used the on_train_epoch_end callback to set those parameters to requires_grad=False. The issue I encounter now is that immediately after that, it hits a DDP unused parameters issue. Is there a correct way of freezing a portion of the model after n steps and allowing continuous training?