Yes, totally agree. These implementations are meant for educational purposes. You could in theory use them to train a model though (GPT-2 also had a from-scratch implementation if I recall correctly). In practice, you probably want to use FlashAttention though. You use it through `torch.nn.functional.scaled_dot_product_attention` etc.