construct_relevancy_map() computation quite slow

Hi thanks for the great work and the demo.ipynb!

Unfortunately while the notebook shows about 1.38s/steps for the relevancy map computation, I am getting quite slower (170s/step) times on google colab. I haven't changed any of the code.

The objects are all on cpu as loading all the attention maps on gpu would lead to OOM.

Do you have any guesses what's happening and why there's such a big difference? Any suggestions would be extremely helpful!

Thanks!