How are you running the optimizer?
After you call add_neurons, do the layers show up in your model’s state_dict / does the weight parameter show up in model.parameters() ?
It might be that you need to re-register the buffers, but I’m not entirely sure.