Fine-tuning only new weights

I am trying to train a transformer model after extending its vocabulary. The problem is I want to keep the original weights frozen and train only the weights associated with the new vocabulary. I was thinking of doing something like this:

processor = processor() # Loading processor
model = model() # Loading model
for param in model.parameters():
    param.requires_grad = False


Is this a valid option? If not, what other options do I have?

It’s unclear what exactly happens in your code as neither object definitions are posted nor did you explain what the custom methods perform.
If you want to freeze a weight parameter partially, you could zero out its gradients before calling optimizer.step() (which should work for stateless optimizers) or you could restore the original weight values after the update (which would also work for optimizers using running stats).