Thanks for the information.



 Remorax:

Are there optimizations pytorch makes on embeddings even when requires_grad is set to False ?


No, that shouldn’t be the case.
As you’ve described, you are only changing the way the work embedding is calculated and passed to the model, so for debugging you could compare the embedding vectors for a fixed batch and check, if both approaches yield the same output.