Try switching all the .cuda() to .mps() I got a 100x speedup on a different language model on a Macbook M1 Air.

https://pytorch.org/docs/stable/notes/mps.html