Why not use nn.MultiheadAttention in vit?

[ZhiyuanChen]

It seems PyTorch have provided torch.nn.MultiheadAttention like flax.nn.base.MultiHeadDotProductAttention, why not use this but write an own Attention?

[rwightman]

@ZhiyuanChen I wasn't quite sure how the official version would look wrt to the attention module and how close it'd be to the PyTorch impl when I started. Plus it was pretty straightforward to just implement it as it is. I don't think the current PyTorch impl s much faster. The Apex one would likely be, but it's harder to work with.

[ZhiyuanChen]

Thank you