[BugFix] remove update from add when transform is not in-place to keep new metadata (#3050)

Author: vmoens

torchrl/data/replay_buffers/replay_buffers.py

@@ -705,15 +705,16 @@ def add(self, data: Any) -> int:
                 make_none = False
                 # Transforms usually expect a time batch dimension when called within a RB, so we unsqueeze the data temporarily
                 is_tc = is_tensor_collection(data)
-                with data.unsqueeze(-1) if is_tc else contextlib.nullcontext(
-                    data
-                ) as data_unsq:
+                cm = data.unsqueeze(-1) if is_tc else contextlib.nullcontext(data)
+                new_data = None
+                with cm as data_unsq:
                     data_unsq_r = self._transform.inv(data_unsq)
                     if is_tc and data_unsq_r is not None:
                         # this is a no-op whenever the result matches the input
-                        data_unsq.update(data_unsq_r)
+                        new_data = data_unsq_r.squeeze(-1)
                     else:
                         make_none = data_unsq_r is None
+                data = new_data if new_data is not None else data
                 if make_none:
                     data = None
         if data is None:

torchrl/envs/transforms/transforms.py

@@ -4507,9 +4507,9 @@ class DeviceCastTransform(Transform):
     """Moves data from one device to another.
 
     Args:
-        device (torch.device or equivalent): the destination device.
-        orig_device (torch.device or equivalent): the origin device. If not specified and
-            a parent environment exists, it it retrieved from it. In all other cases,
+        device (torch.device or equivalent): the destination device (outside the environment or buffer).
+        orig_device (torch.device or equivalent): the origin device (inside the environment or buffer).
+            If not specified and a parent environment exists, it it retrieved from it. In all other cases,
             it remains unspecified.
 
     Keyword Args: