Question about TorchRL ParallelEnv error on single-gpu device

Ivar_Gaitan: <p>Hi! When creating a simple <code>ParallelEnv</code> on a single CUDA device, I get a runtime error. While it may not be a <a href="https://discuss.pytorch.org/t/using-cuda-multiprocessing-with-single-gpu/7300/4">recommended</a> practice, I don’t understand why it fails. Here’s the code (modified from the TorchRL tutorials):</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torchrl.envs import GymEnv, ParallelEnv


def main():
    def env_make():
        return GymEnv("Pendulum-v1", device=torch.device("cuda:0"))

    parallel_env = ParallelEnv(1, env_make)
    parallel_env.reset()
    parallel_env.close()


if __name__ == "__main__":
    main()

</code></pre>
<p>Om my single-GPU device (RTX 4060) this results in the following error:</p>
<pre data-code-wrap="python"><code class="lang-python">Traceback (most recent call last):
  File "/.../.local/share/uv/python/cpython-3.11.11-linux-x86_64-gnu/lib/python3.11/multiprocessing/spawn.py", line 122, in spawn_main
  File "/.../torchrl-learn/test.py", line 16, in &lt;module&gt;
    main()
    exitcode = _main(fd, parent_sentinel)
  File "/.../torchrl-learn/test.py", line 10, in main
               ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../.local/share/uv/python/cpython-3.11.11-linux-x86_64-gnu/lib/python3.11/multiprocessing/spawn.py", line 132, in _main
    parallel_env.reset()
    self = reduction.pickle.load(from_parent)
  File "/.../torchrl-learn/.venv/lib/python3.11/site-packages/torchrl/envs/common.py", line 2858, in reset
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../torchrl-learn/.venv/lib/python3.11/site-packages/torch/multiprocessing/reductions.py", line 181, in rebuild_cuda_tensor
    storage = storage_cls._new_shared_cuda(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../torchrl-learn/.venv/lib/python3.11/site-packages/torch/storage.py", line 1452, in _new_shared_cuda
    tensordict_reset = self._reset(tensordict, **kwargs)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../torchrl-learn/.venv/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return torch.UntypedStorage._new_shared_cuda(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: invalid resource handle


</code></pre>

Ivar_Gaitan: <p>Thanks for the response. However, I’m not able to make it work with your suggested fix. Moving the device object out of the parameter list as such does not produce any different result.</p>
<pre data-code-wrap="python"><code class="lang-python">def env_make():
    device = torch.device("cuda:0")
    return GymEnv("Pendulum-v1", device=device)
</code></pre>
<p>Also, I haven’t found anything corresponding to <code>VecEnv</code> in TorchRL, could you please elaborate what you mean?</p>

vmoens: <p>There is no <code>VecEnv</code> in torchrl. Also the code snippet you shared works on other cuda architecture so I’m not sure what is going on here.<br>
Can you share:</p>
<ul>
<li>pytorch version</li>
<li>torchrl version</li>
<li>cuda version</li>
</ul>

Ivar_Gaitan: <p>Absolutely,</p>
<p>CUDA: 12.9<br>
pytorch: 2.7.1<br>
torchrl: 0.9.2</p>