RETRO: RuntimeError: stack expects each tensor to be equal size, but got [2, 32] at entry 0 and [1, 32] at entry 29

Hi,

Running the exact code on github for deepmind's retrieval transformer - RETRO, getting the following error:

RuntimeError: stack expects each tensor to be equal size, but got [2, 32] at entry 0 and [1, 32] at entry 29

Could you please help me with this? I used the same dataset as in the code.