[Deprecation] Deprecate ambiguous device for memmap replay buffer (#1624)

Author: Vincent Moens

torchrl/data/replay_buffers/storages.py

@@ -586,8 +586,14 @@ def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
                 data.clone()
                 .expand(self.max_size, *data.shape)
                 .memmap_like(prefix=self.scratch_dir)
-                .to(self.device)
             )
+            if self.device.type != "cpu":
+                warnings.warn(
+                    "Support for Memmap device other than CPU will be deprecated in v0.4.0.",
+                    category=DeprecationWarning,
+                )
+                out = out.to(self.device).memmap_()
+
             for key, tensor in sorted(
                 out.items(include_nested=True, leaves_only=True), key=str
             ):
@@ -603,8 +609,14 @@ def _init(self, data: Union[TensorDictBase, torch.Tensor]) -> None:
                 data.clone()
                 .expand(self.max_size, *data.shape)
                 .memmap_like(prefix=self.scratch_dir)
-                .to(self.device)
             )
+            if self.device.type != "cpu":
+                warnings.warn(
+                    "Support for Memmap device other than CPU will be deprecated in v0.4.0.",
+                    category=DeprecationWarning,
+                )
+                out = out.to(self.device).memmap_()
+
             for key, tensor in sorted(
                 out.items(include_nested=True, leaves_only=True), key=str
             ):

tutorials/sphinx-tutorials/pretrained_models.py

@@ -88,7 +88,7 @@
 #
 from torchrl.data import LazyMemmapStorage, ReplayBuffer
 
-storage = LazyMemmapStorage(1000, device=device)
+storage = LazyMemmapStorage(1000)
 rb = ReplayBuffer(storage=storage, transform=r3m)
 
 ##############################################################################