I compared your code against lucidrains'.
Yours:
Line 330 in 2ae8437
Lucidrains':
https://github.com/lucidrains/vit-pytorch/blob/9992a615d1f3e0117a62ced39a88d48b6be68858/vit_pytorch/na_vit.py#L318
I doubt there might be a bug at this line in your code.
Pay now to fund the work behind this issue.
Get updates on progress being made.
Maintainer is rewarded once the issue is completed.
You're funding impactful open source efforts
You want to contribute to this effort
You want to get funding like this too