Your code in the first block looks like it would work. I suspect the problem is one of two things:

The dropout may prevent the model from fully converging.
The default learning rate for Adam is 1e-3, which is much too high for training Transformer models. Try learning rates in the range 1e-5 to 1e-4. If training loss is still not decreasing even with a lower learning rate and no dropout then let me know and I’ll investigate further.