modeling_fnet.py config option 
https://github.com/huggingface/transformers/blob/b7d264be0d70d5309dbf24a6c7ebba1e073bdda5/src/transformers/models/fnet/modeling_fnet.py#L177

Should refer to config.tpu_short_seq_length

Maybe of interest to @gchhablani :)<|||||>So sorry about that @ontocord.

I will fix it immediately. Thanks a lot for pointing it out.<|||||>I made a few more fixes. See this notebook. Basically I moved the buffer for doing fft to the encoder class, and made it a parameter of the module. in this way you can fix problems with cuda vs non cuda. the parameter weren't converting to cuda. i also made one fix relating to dtype conversion so you can run in half mode. You can do a diff and feel free to accept and PR it if you wish!

https://colab.research.google.com/drive/19A-qrta2yqVcfSjUtjEv6w1HqOyYXuEo#scrollTo=olgITNUmyQxA
<|||||>As an aside @gchhablani, did you run any speed tests. My small test didn't show much diff w/ bert. but I suppose the authors said that training is faster. I supposed I have to test it on something longer?

https://arxiv.org/abs/2105.03824<|||||>@ontocord I ran fine-tuning script on GLUE and the FNet model is faster than BERT. It takes 70% time (train + eval). The original ratios mentioned are only training and model is in Flax.

If you are talking about fourier transform only, then I didn't perform a test of speed for the three different ways they are implemented - GPU, TPU short sequence, TPU long sequence. I assume it will be slower because of PyTorch.

An interesting comparison could be creating a HuggingFace Flax model and comparing with the original. 

@patrickvonplaten what do you think?<|||||>Could you all share the fine-tuning script or a colab that shows the run? 70% of the time of BERT or 70% faster than BERT?
<|||||>@ontocord It is 70% of the time of BERT in our case, but our script is different, and we used PyTorch. Torch does not have `vmap` as of now, because of which I expect slow down in the performance.

You can see the runs (go to Metrics) in the model cards: https://huggingface.co/models?other=fnet-bert-base-comparison<|||||>@ontocord I checked your changes, maybe you can open a PR and we can discuss? It'll be easier to involve others as well.<|||||>Sorry didn't see your message. Looks like you already did this... Great job :)