Revert "[BugFix] Fix tensordict private imports" (#2276)

pytorch · Jul 8, 2024 · a151923 · a151923
1 parent 95b6a16
commit a151923
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/test/test_cost.py b/test/test_cost.py
@@ -12,7 +12,7 @@
 from dataclasses import asdict, dataclass
 
 from packaging import version as pack_version
-from tensordict._C import unravel_keys
+from tensordict._tensordict import unravel_keys
 
 from tensordict.nn import (
     InteractionType,

diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -27,7 +27,7 @@
     TensorDict,
     TensorDictBase,
 )
-from tensordict._C import unravel_key
+from tensordict._tensordict import unravel_key
 from torch import multiprocessing as mp
 from torchrl._utils import (
     _check_for_faulty_process,

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -39,7 +39,7 @@
     unravel_key,
     unravel_key_list,
 )
-from tensordict._C import _unravel_key_to_tuple
+from tensordict._tensordict import _unravel_key_to_tuple
 from tensordict.nn import dispatch, TensorDictModuleBase
 from tensordict.utils import expand_as_right, expand_right, NestedKey
 from torch import nn, Tensor