Suport device=None (#1025)

Co-authored-by: Molly Xu <[email protected]>

Author: mollyxu

src/torchcodec/decoders/_video_decoder.py

@@ -55,7 +55,8 @@ class VideoDecoder:
             decoding which is best if you are running a single instance of ``VideoDecoder``.
             Passing 0 lets FFmpeg decide on the number of threads.
             Default: 1.
-        device (str or torch.device, optional): The device to use for decoding. Default: "cpu".
+        device (str or torch.device, optional): The device to use for decoding.
+            If ``None`` (default), uses the current default device.
             If you pass a CUDA device, we recommend trying the "beta" CUDA
             backend which is faster! See :func:`~torchcodec.decoders.set_cuda_backend`.
         seek_mode (str, optional): Determines if frame access will be "exact" or
@@ -102,7 +103,7 @@ def __init__(
         stream_index: Optional[int] = None,
         dimension_order: Literal["NCHW", "NHWC"] = "NCHW",
         num_ffmpeg_threads: int = 1,
-        device: Optional[Union[str, torch_device]] = "cpu",
+        device: Optional[Union[str, torch_device]] = None,
         seek_mode: Literal["exact", "approximate"] = "exact",
         custom_frame_mappings: Optional[
             Union[str, bytes, io.RawIOBase, io.BufferedReader]
@@ -143,7 +144,9 @@ def __init__(
         if num_ffmpeg_threads is None:
             raise ValueError(f"{num_ffmpeg_threads = } should be an int.")
 
-        if isinstance(device, torch_device):
+        if device is None:
+            device = str(torch.get_default_device())
+        elif isinstance(device, torch_device):
             device = str(device)
 
         device_variant = _get_cuda_backend()

test/test_decoders.py

@@ -388,6 +388,31 @@ def test_device_instance(self):
         decoder = VideoDecoder(NASA_VIDEO.path, device=torch.device("cpu"))
         assert isinstance(decoder.metadata, VideoStreamMetadata)
 
+    @pytest.mark.parametrize(
+        "device_str",
+        [
+            "cpu",
+            pytest.param("cuda", marks=pytest.mark.needs_cuda),
+        ],
+    )
+    def test_device_none_default_device(self, device_str):
+        # VideoDecoder defaults to device=None, which should respect both
+        # torch.device() context manager and torch.set_default_device().
+
+        # Test with context manager
+        with torch.device(device_str):
+            decoder = VideoDecoder(NASA_VIDEO.path)
+            assert decoder[0].device.type == device_str
+
+        # Test with set_default_device
+        original_device = torch.get_default_device()
+        try:
+            torch.set_default_device(device_str)
+            decoder = VideoDecoder(NASA_VIDEO.path)
+            assert decoder[0].device.type == device_str
+        finally:
+            torch.set_default_device(original_device)
+
     @pytest.mark.parametrize("device", all_supported_devices())
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_getitem_fails(self, device, seek_mode):