The speed.

[wangyuxin87]

Thanks for your excellent work.
However, GAU is slower than the original MHSA in my implementation, 3.5s vs 0.7s. As I simply use "from flash_pytorch import GAU" with the default setting.
I there something wrong with my implementation?