Cannot import name 'LSTM' from 'torchrl.modules.tensordict_module.rnn'

1

So after about 8 months of wishing I could use VMap to stack LSTMs, I’ve finally found this: LSTM — torchrl main documentation

On this page, I see: “This class is implemented without relying on CuDNN, which makes it compatible with torch.vmap()”, which is great for me. However, when following the instructions on that page, I run into an error at the very first step, the import:

from torchrl.modules.tensordict_module.rnn import LSTM

That line produces this error:

ImportError: cannot import name 'LSTM' from 'torchrl.modules.tensordict_module.rnn' (/home/username/miniconda3/envs/envname/lib/python3.8/site-packages/torchrl/modules/tensordict_module/rnn.py)

Versioning info for context:

>>> import torch
>>> torch.__version__
>>> '2.1.1'
>>> import torchrl
>>> torchrl.__version__
>>> '0.2.1'

And python itself is 3.8.18.

TorchRL was installed via:

pip install torchrl

as per instructions on this page: Introduction to TorchRL — torchrl main documentation

Any takers on this? Frustrating error because I’ve been waiting for a way to stack LSTMs for almost a year now, and I finally find this torchrl LSTM that can supposedly be VMapped, seemingly follow usage instructions exactly, and can’t get past the import step.

Mega thanks to anyone that can help with this (:

2

Ok nevermind, problem solved. For posterity: tensordict_module.rnn.LSTM only currently ( as of 27.01.2024 ) exists in the nightly build of torchrl, which requires one to have the nightly build of pytorch itself installed. Uninstalling torch 2.1.1, then installing the nightly build of pytorch, then installing the nightly build of torchrl ( pytorch nightly installed via conda instructions from this page and torchrl nightly installed via pip install torchrl-nightly ) worked to solve this problem for me.

3

Yes this feature will be in the v0.3 release next week. It should work with torch.compile as well as vmap.

1 Like