arvindmohan:

I ended up implementing LSTM from scratch in Pytorch as it gave me greater flexibility to try various modifications and not mess with the Pytorch internals. It didn’t take much time though.


Thank you Arvind!
It’s a great solution, will try it out soon!
I need to initialise hidden states querying a big tensor that I would avoid coping to each gpu.
Although pretty inefficient I think I have to send the query vector to where the big tensor is, extracting the query and then call .cuda() again. Not sure if autograd can track these device changes… hope to give a feedback soon