Reformer model with long sequence throws an error

[satya77]

Im trying to train the basic Reformer and not the RefomerLM on long sequence of text based on the language generation example. Simply by replacing the RefomerLM class with the reformer and remove the mask, but feeding in the entire crime and punishment book, throws the following error:

TypeError: requesting more random bits than a single call provides.

everything works fine if I cut down the input to smaller sequences. The example can be seen in the following notebook:
https://colab.research.google.com/drive/1C9KOHOfuVhoOqzRKx_rRaeOV3jPvuZ_L

[lukaszkaiser]

This looks like a potential JAX problem. Could you paste here the whole error log?

[lukaszkaiser]

Closing as it doesn't look active, please reopen if needed of course!