remove _validate_cvcuda_dtype check

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -6747,63 +6747,28 @@ def test_functional_error(self):
 class TestToCVCUDATensor:
     """Tests for to_cvcuda_tensor function following patterns from TestToPil"""
 
-    def test_1_channel_uint8_tensor_to_cvcuda_tensor(self):
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64])
+    def test_1_channel_to_cvcuda_tensor(self, dtype):
         # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.randint(0, 256, (1, 4, 4), dtype=torch.uint8)
-        img_data = img_data.cuda()
-        cvcuda_img = F.to_cvcuda_tensor(img_data)
-        # Check that the conversion succeeded and format is correct
-        assert cvcuda_img is not None
-
-    def test_1_channel_int16_tensor_to_cvcuda_tensor(self):
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.randint(0, 256, (1, 4, 4), dtype=torch.int16)
-        img_data = img_data.cuda()
-        cvcuda_img = F.to_cvcuda_tensor(img_data)
-        assert cvcuda_img is not None
-
-    def test_1_channel_int32_tensor_to_cvcuda_tensor(self):
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.randint(0, 256, (1, 4, 4), dtype=torch.int32)
-        img_data = img_data.cuda()
-        cvcuda_img = F.to_cvcuda_tensor(img_data)
-        assert cvcuda_img is not None
-
-    def test_1_channel_float32_tensor_to_cvcuda_tensor(self):
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.rand(1, 4, 4)
-        img_data = img_data.cuda()
-        cvcuda_img = F.to_cvcuda_tensor(img_data)
-        assert cvcuda_img is not None
-
-    def test_3_channel_uint8_tensor_to_cvcuda_tensor(self):
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.randint(0, 256, (3, 4, 4), dtype=torch.uint8)
+        if dtype in (torch.uint8, torch.uint16):
+            img_data = torch.randint(0, 256, (1, 4, 4), dtype=dtype)
+        else:
+            img_data = torch.rand(1, 4, 4, dtype=dtype)
         img_data = img_data.cuda()
         cvcuda_img = F.to_cvcuda_tensor(img_data)
         assert cvcuda_img is not None
 
-    def test_3_channel_float32_tensor_to_cvcuda_tensor(self):
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64])
+    def test_3_channel_to_cvcuda_tensor(self, dtype):
         # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.rand(3, 4, 4)
+        if dtype in (torch.uint8, torch.uint16):
+            img_data = torch.randint(0, 256, (3, 4, 4), dtype=dtype)
+        else:
+            img_data = torch.rand(3, 4, 4, dtype=dtype)
         img_data = img_data.cuda()
         cvcuda_img = F.to_cvcuda_tensor(img_data)
         assert cvcuda_img is not None
 
-    def test_unsupported_num_channels(self):
-        # Test 2-channel image (not supported)
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.rand(2, 5, 5)
-        img_data = img_data.cuda()
-        with pytest.raises(ValueError, match="Only 1 and 3 channel images are supported"):
-            F.to_cvcuda_tensor(img_data)
-
-        # Test 4-channel image (not supported)
-        img_data = torch.randint(0, 256, (4, 5, 5), dtype=torch.uint8)
-        img_data = img_data.cuda()
-        with pytest.raises(ValueError, match="Only 1 and 3 channel images are supported"):
-            F.to_cvcuda_tensor(img_data)
-
     def test_invalid_input_type(self):
         with pytest.raises(TypeError, match=r"pic should be `torch.Tensor`"):
             F.to_cvcuda_tensor("invalid_input")
@@ -6828,32 +6793,12 @@ def test_invalid_dimensions(self):
             img_data = img_data.cuda()
             F.to_cvcuda_tensor(img_data)
 
-    def test_float64_tensor_to_cvcuda_tensor(self):
-        # Test single channel float64 (F64 format is supported)
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.rand(1, 4, 4, dtype=torch.float64)
-        img_data = img_data.cuda()
-        cvcuda_img = F.to_cvcuda_tensor(img_data)
-        assert cvcuda_img is not None
-
-    def test_float64_rgb_not_supported(self):
-        # Test 3-channel float64 is NOT supported (no RGBf64 format in CV-CUDA)
-        # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        img_data = torch.rand(3, 4, 4, dtype=torch.float64)
-        img_data = img_data.cuda()
-        with pytest.raises(TypeError, match=r"Unsupported dtype"):
-            F.to_cvcuda_tensor(img_data)
-
     @pytest.mark.parametrize("num_channels", [1, 3])
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.int32, torch.float32, torch.float64])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64])
     def test_round_trip(self, num_channels, dtype):
-        # Skip int16/int32/float64 for 3-channel (only supported for single-channel)
-        if num_channels == 3 and dtype in (torch.int16, torch.int32, torch.float64):
-            pytest.skip(f"{dtype} is only supported for single-channel images")
-
         # Setup: Create a tensor in CHW format (PyTorch standard)
         # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        if dtype in (torch.uint8, torch.int16, torch.int32):
+        if dtype in (torch.uint8, torch.uint16):
             original_tensor = torch.randint(0, 256, (num_channels, 4, 4), dtype=dtype)
         else:
             original_tensor = torch.rand(num_channels, 4, 4, dtype=dtype)
@@ -6871,16 +6816,12 @@ def test_round_trip(self, num_channels, dtype):
         torch.testing.assert_close(result_tensor, original_tensor, rtol=0, atol=0)
 
     @pytest.mark.parametrize("num_channels", [1, 3])
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.int32, torch.float32, torch.float64])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64])
     @pytest.mark.parametrize("batch_size", [1, 2, 4])
     def test_round_trip_batched(self, num_channels, dtype, batch_size):
-        # Skip int16/int32/float64 for 3-channel (only supported for single-channel)
-        if num_channels == 3 and dtype in (torch.int16, torch.int32, torch.float64):
-            pytest.skip(f"{dtype} is only supported for single-channel images")
-
         # Setup: Create a batched tensor in NCHW format
         # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        if dtype in (torch.uint8, torch.int16, torch.int32):
+        if dtype in (torch.uint8, torch.uint16):
             original_tensor = torch.randint(0, 256, (batch_size, num_channels, 4, 4), dtype=dtype)
         else:
             original_tensor = torch.rand(batch_size, num_channels, 4, 4, dtype=dtype)

torchvision/transforms/v2/functional/_type_conversion.py

@@ -53,34 +53,6 @@ def to_image(inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray]) -> tv_tenso
 pil_to_tensor = _F.pil_to_tensor
 
 
-def _validate_cvcuda_dtype(img_tensor: torch.Tensor) -> None:
-    """Validate that tensor dtype and channel count are supported by CV-CUDA.
-
-    Args:
-        img_tensor: Tensor with shape (H, W, C) where C is number of channels.
-
-    Raises:
-        TypeError: If dtype is not supported for the given number of channels.
-        ValueError: If number of channels is not 1 or 3.
-    """
-    num_channels = img_tensor.shape[2]
-    dtype = img_tensor.dtype
-
-    # Handle single channel images
-    if num_channels == 1:
-        if dtype not in (torch.uint8, torch.int16, torch.int32, torch.float32, torch.float64):
-            raise TypeError(f"Unsupported dtype {dtype} for single channel image")
-
-    # Handle 3 channel images (defaults to RGB)
-    elif num_channels == 3:
-        if dtype not in (torch.uint8, torch.float32):
-            # Note: CV-CUDA does not support float64 for RGB images (only F64 for single-channel)
-            raise TypeError(f"Unsupported dtype {dtype} for 3-channel image")
-
-    else:
-        raise ValueError(f"Only 1 and 3 channel images are supported. Got {num_channels} channels.")
-
-
 @torch.jit.unused
 def to_cvcuda_tensor(pic) -> "cvcuda.Tensor":
     """Convert a torch.Tensor to cvcuda.Tensor. This function does not support torchscript.
@@ -113,14 +85,9 @@ def to_cvcuda_tensor(pic) -> "cvcuda.Tensor":
     else:
         raise ValueError(f"pic should be 3 or 4 dimensional. Got {pic.ndim} dimensions.")
 
-    # At this point, img_tensor is always 4D in NCHW format
     # Convert NCHW -> NHWC
     img_tensor = img_tensor.permute(0, 2, 3, 1)
 
-    # Validate dtype and channel count from the first image
-    sample_img = img_tensor[0]
-    _validate_cvcuda_dtype(sample_img)
-
     # Convert to CV-CUDA tensor with NHWC layout
     return cvcuda.as_tensor(img_tensor.cuda().contiguous(), cvcuda.TensorLayout.NHWC)