[Feature] VideoClipRef.decode: honor the container device (#3839)

Author: vmoens

test/transforms/test_video_transforms.py

@@ -572,6 +572,18 @@ def test_transform_device(self, video_path):
         )(td)
         assert out["pixels"].device.type == "cuda"
 
+    def test_tensordict_to_device_decodes_on_device(self, video_path):
+        # Moving the tensordict that holds the reference makes decode() materialize
+        # frames on that device (decode falls back to the container device).
+        td = TensorDict({"frame": VideoClipRef.from_file(video_path)}, batch_size=[20])
+        td = td.to("cpu")
+        assert td["frame"].decode().device.type == "cpu"
+        td = td.to("cuda")
+        assert td["frame"].decode().device.type == "cuda"
+        assert td["frame"][3:7].decode().device.type == "cuda"
+        # an explicit override still wins
+        assert td["frame"].decode(device="cpu").device.type == "cpu"
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()

torchrl/data/video.py

@@ -726,9 +726,11 @@ def decode(self, *, device: Any = None, dtype: Any = None) -> torch.Tensor:
 
         Keyword Args:
             device (torch.device or str, optional): output device for the decoded
-                frames, overriding ``out_device``. A CUDA device uses GPU (NVDEC)
-                decoding when the torchcodec build supports it, and otherwise decodes
-                on CPU and moves the frames to the device.
+                frames, overriding ``out_device``. Defaults to ``out_device`` if set,
+                else the reference's own device (so ``td.to("cuda")`` on a tensordict
+                holding the reference decodes onto CUDA), else CPU. A CUDA device uses
+                GPU (NVDEC) decoding when the torchcodec build supports it, and
+                otherwise decodes on CPU and moves the frames to the device.
             dtype (torch.dtype, optional): dtype for the decoded frames, overriding
                 ``out_dtype``. Defaults to ``uint8``.
 
@@ -739,6 +741,11 @@ def decode(self, *, device: Any = None, dtype: Any = None) -> torch.Tensor:
             raise ModuleNotFoundError(_TORCHCODEC_ERROR)
         if device is None:
             device = _first(self.out_device)
+        if device is None:
+            # Fall back to the reference's own (container) device, so moving the
+            # tensordict that holds it -- e.g. ``td.to("cuda")`` -- makes
+            # ``decode()`` materialize frames on that device.
+            device = self.device
         if dtype is None:
             dtype = _first(self.out_dtype)
         stream = _first(self.stream)