search

pytorch-labs / LeanRL

LeanRL is a fork of CleanRL, where selected PyTorch scripts optimized for performance using compile and cudagraphs.
Other
354 stars 9 forks source link

Can't import from_module from tensordict #5

Open StoneT2000 opened 6 days ago

StoneT2000 commented 6 days ago

Hi, thanks for the awesome work in making RL with torch fast.

The general lessons (e.g. how to use compile) are making some improvements but they are very marginal. I want to try and use the cudagraphs improvements which bring the biggest speedups, but am having some issues with the from_module function, it seems i can't import it?

  File "/home/stao/work/research/maniskill/ManiSkill/examples/baselines/ppo/ppo.py", line 23, in <module>
    from tensordict import from_module
ImportError: cannot import name 'from_module' from 'tensordict' (/home/stao/miniforge3/envs/ms3-py311/lib/python3.11/site-packages/tensordict/__init__.py)

This is my current used version for torch and tensordict

Name: tensordict-nightly
Version: 2024.7.3
Summary: 
Home-page: https://github.com/pytorch/tensordict
Author: tensordict contributors
Author-email: vmoens@fb.com
License: BSD
Name: torch
Version: 2.4.1
Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration
Home-page: https://pytorch.org/
Author: PyTorch Team
Author-email: packages@pytorch.org
License: BSD-3

Any idea?

I am also trying to do as minimal of an install as possible as I want to make some torch compile examples in a separate project, so I hope there isn't some dependency I missed. I only installed tensordict and torch essentially.

vmoens commented 6 days ago

It's a feature on the most recent nightly, I think that here the issue is that you have installed the latest nightly that is compatible with torch 2.4 - try using torch nightly too! Sorry for the inconvenience, I'll upgrade it all to non-nightly asap