[Versioning] v0.6.1 fixes

benchmarks/test_replaybuffer_benchmark.py

@@ -13,6 +13,7 @@
     LazyMemmapStorage,
     LazyTensorStorage,
     ListStorage,
+    ReplayBuffer,
     TensorDictPrioritizedReplayBuffer,
     TensorDictReplayBuffer,
 )

setup.py

@@ -176,7 +176,7 @@ def _main(argv):
     if is_nightly:
         tensordict_dep = "tensordict-nightly"
     else:
-        tensordict_dep = "tensordict>=0.6.1"
+        tensordict_dep = "tensordict>=0.6.2"
 
     if is_nightly:
         version = get_nightly_version()

test/test_rb.py

@@ -26,8 +26,14 @@
     )
     from pytorch.rl.test.mocking_classes import CountingEnv
 else:
-    from _utils_internal import CARTPOLE_VERSIONED, get_default_devices, make_tc
+    from _utils_internal import (
+        capture_log_records,
+        CARTPOLE_VERSIONED,
+        get_default_devices,
+        make_tc,
+    )
     from mocking_classes import CountingEnv
+
 from packaging import version
 from packaging.version import parse
 from tensordict import (
@@ -119,6 +125,7 @@
 _has_gym = importlib.util.find_spec("gym") is not None
 _has_snapshot = importlib.util.find_spec("torchsnapshot") is not None
 _os_is_windows = sys.platform == "win32"
+TORCH_VERSION = version.parse(version.parse(torch.__version__).base_version)
 
 torch_2_3 = version.parse(
     ".".join([str(s) for s in version.parse(str(torch.__version__)).release])

test/test_utils.py

@@ -15,7 +15,7 @@
 import torch
 
 if os.getenv("PYTORCH_TEST_FBCODE"):
-    from pytorch.rl.test._utils_internal import capture_log_records, get_default_devices
+    from pytorch.rl.test._utils_internal import get_default_devices
 else:
     from _utils_internal import get_default_devices
 from torchrl._utils import _rng_decorator, get_binary_env_var, implement_for