Request for help for LSHSelfAttention()

Hi @lucidrains thank you for your excellent work (I star it).

I am trying to use the LSHSelfAttention() layer in my network instead of my transformer encoder layer.

A pseudocode of what I am doing is that:

word_embeddings = word_embeddings(input)  # batch, seq_len, emb_dim
lsh_encoded = self.lsh_self_attention(word_embeddings)

I continuously get a vector of NaN values, to avoid it I decrease my learning rate from 1e-3 to 1e-5, but nothing is changed.

1. Am I using the correct layer?
2. Should I use Reformer() instead of LSHSelfAttention()? I tried to use Reformer() but I also get an error there, which tells me that my sequence is divisible by the number of buckets (I’m still working on it).

[lucidrains]

Woohoo! Congrats 💯

[andreabac3]

@lucidrains I am using pytorch_lightning==0.8.5 so I suppose, looking in the default parameter in Trainer() class, it’s set as O2

amp_level: str = 'O2',   # backward compatible, todo: remove in v1.0.0

which is located in $HOME/.local/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py

PS: It’s works with lr = 1e-3 without problem.