[torchcodec] Allow sampler to use GPU decoding (meta-pytorch#136)

Summary: Pull Request resolved: meta-pytorch#136

Reviewed By: scotts

Differential Revision: D60402239

Author: ahmadsharif1

src/torchcodec/_samplers/video_clip_sampler.py

@@ -31,6 +31,7 @@ class VideoTooShortException(Exception):
 @dataclass
 class DecoderArgs:
     num_threads: int = 0
+    device: torch.device = torch.device("cpu")
 
 
 @dataclass
@@ -163,6 +164,7 @@ def forward(self, video_data: Tensor) -> Union[List[Any]]:
             width=target_width,
             height=target_height,
             num_threads=self.decoder_args.num_threads,
+            device_string=str(self.decoder_args.device),
         )
 
         clips: List[Any] = []

test/samplers/test_video_clip_sampler.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 from torchcodec._samplers import (
+    DecoderArgs,
     IndexBasedSamplerArgs,
     TimeBasedSamplerArgs,
     VideoArgs,
@@ -30,11 +31,16 @@
         ),
     ],
 )
-def test_sampler(sampler_args):
+@pytest.mark.parametrize(("device"), [torch.device("cpu"), torch.device("cuda:0")])
+def test_sampler(sampler_args, device):
+    if device.type == "cuda" and not torch.cuda.is_available():
+        pytest.skip("GPU not available")
+
     torch.manual_seed(0)
     desired_width, desired_height = 320, 240
     video_args = VideoArgs(desired_width=desired_width, desired_height=desired_height)
-    sampler = VideoClipSampler(video_args, sampler_args)
+    decoder_args = DecoderArgs(device=device)
+    sampler = VideoClipSampler(video_args, sampler_args, decoder_args)
     clips = sampler(NASA_VIDEO.to_tensor())
     assert_tensor_equal(len(clips), sampler_args.clips_per_video)
     clip = clips[0]