Great! I’ll check it out. I did some changes to the way I wrote the functions in Python and I got better performance than LSTM with JIT version on CPU, by lot. But when I use GPU the difference is massive in favor of LSTM. So I guess I’ll have to code it by hand. It is time consuming but in the end I think it’s something I should know how to do.
It’s tricky because I tried using addmm instead of F.linear and pre-transposing the weights before the time loop (which should give better performance) in my Python code,  but on GPU I get almost no improvements.
Thanks for all the help.
Best
PS: I noticed that LSTMs use two bias vectors instead of one, and I know that this is done for keeping it similar to cuDNN but I can’t find the reason why they defined the model in that way. What’s more interesting is that in the link you posted earlier, it seems to imply only one bias. Any ideas on this?
Just read that you don’t know why this is on your PR.