Error when mask is of the same size as that of the input X

[Vaibhavs10]

Hi,

First of all, thank you for putting such an easy to use implementation on GitHub.
I'm trying to incorporate the nystrom attention into a legacy codebase, it previously used to provide the input X and the mask (off the same dimensions as X) to a Multi headed Attention Layer.

When I'm trying to integrate nystrom attention with it, it runs alright without the mask.
But, when I pass the mask alongside it, it throws einops rearrange error.

Sorry, if this is a very basic question, but how would you recommend I deal with handling 3D mask (same dimensions as the size of input) in the codebase.

Best,
VB

[Vaibhavs10]

I got confused with the dimensions, this is fixed now