Reformer for Multi-GPU not Possible for Torch > 1.4.0

As mention in this issue, it is not possible to train or apply Reformer (to create text representations of new data) with torch > 1.4.0. This makes it difficult to use the Reformer model architecture in any environment using recent releases of PyTorch (latest is 1.9.0). Would anyone be interested in helping to resolve this?

I’d be happy to help, but I don’t understand the code well enough to spearhead the fix. I’m also aware of a few other users who are stuck using an old version of PyTorch so that they can use Reformer or who are actively avoiding Reformer so that they can keep their PyTorch version current.