fix: to_dtype_cvcuda conventions

Author: justincdavis

test/test_transforms_v2.py

@@ -2628,7 +2628,17 @@ def test_kernel(self, kernel, make_input, input_dtype, output_dtype, device, sca
             scale=scale,
         )
 
-    @pytest.mark.parametrize("make_input", [make_image_tensor, make_image, make_video])
+    @pytest.mark.parametrize(
+        "make_input",
+        [
+            make_image_tensor,
+            make_image,
+            make_video,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
+    )
     @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
     @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
     @pytest.mark.parametrize("device", cpu_and_cuda())
@@ -2643,7 +2653,16 @@ def test_functional(self, make_input, input_dtype, output_dtype, device, scale):
 
     @pytest.mark.parametrize(
         "make_input",
-        [make_image_tensor, make_image, make_bounding_boxes, make_segmentation_mask, make_video],
+        [
+            make_image_tensor,
+            make_image,
+            make_bounding_boxes,
+            make_segmentation_mask,
+            make_video,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
     )
     @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
     @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
@@ -2807,38 +2826,12 @@ def test_uint16(self):
         assert_equal(F.to_dtype(img_float32, torch.uint8, scale=True), img_uint8)
         assert_close(F.to_dtype(img_uint8, torch.float32, scale=True), img_float32, rtol=0, atol=1e-2)
 
-
-@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="cvcuda is not available")
-@needs_cuda
-class TestToDtypeCVCUDA:
-    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
-    @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
-    @pytest.mark.parametrize("device", cpu_and_cuda())
-    @pytest.mark.parametrize("scale", (True, False))
-    def test_functional(self, input_dtype, output_dtype, device, scale):
-        check_functional(
-            F.to_dtype,
-            make_image_cvcuda(batch_dims=(1,), dtype=input_dtype, device=device),
-            dtype=output_dtype,
-            scale=scale,
-        )
-
-    @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8])
-    @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8])
-    @pytest.mark.parametrize("device", cpu_and_cuda())
-    @pytest.mark.parametrize("scale", (True, False))
-    @pytest.mark.parametrize("as_dict", (True, False))
-    def test_transform(self, input_dtype, output_dtype, device, scale, as_dict):
-        cvc_input = make_image_cvcuda(batch_dims=(1,), dtype=input_dtype, device=device)
-        if as_dict:
-            output_dtype = {type(cvc_input): output_dtype}
-        check_transform(transforms.ToDtype(dtype=output_dtype, scale=scale), cvc_input, check_sample_input=not as_dict)
-
+    @pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
     @pytest.mark.parametrize("input_dtype", [torch.float32, torch.float64, torch.uint8, torch.uint16])
     @pytest.mark.parametrize("output_dtype", [torch.float32, torch.float64, torch.uint8, torch.uint16])
     @pytest.mark.parametrize("device", cpu_and_cuda())
     @pytest.mark.parametrize("scale", (True, False))
-    def test_image_correctness(self, input_dtype, output_dtype, device, scale):
+    def test_cvcuda_parity(self, input_dtype, output_dtype, device, scale):
         if input_dtype.is_floating_point and output_dtype == torch.int64:
             pytest.xfail("float to int64 conversion is not supported")
         if input_dtype == torch.uint8 and output_dtype == torch.uint16 and device == "cuda":

torchvision/transforms/v2/functional/__init__.py

@@ -158,7 +158,6 @@
     sanitize_bounding_boxes,
     sanitize_keypoints,
     to_dtype,
-    to_dtype_cvcuda,
     to_dtype_image,
     to_dtype_video,
 )

torchvision/transforms/v2/functional/_misc.py

@@ -371,11 +371,13 @@ def _to_dtype_tensor_dispatch(inpt: torch.Tensor, dtype: torch.dtype, scale: boo
     _cvcuda_to_torch_dtypes = {v: k for k, v in _torch_to_cvcuda_dtypes.items()}
 
 
-def to_dtype_cvcuda(
+def _to_dtype_cvcuda(
     inpt: "cvcuda.Tensor",
     dtype: torch.dtype,
     scale: bool = False,
 ) -> "cvcuda.Tensor":
+    cvcuda = _import_cvcuda()
+
     dtype_in = _cvcuda_to_torch_dtypes[inpt.dtype]
     cvc_dtype = _torch_to_cvcuda_dtypes[dtype]
 
@@ -409,7 +411,7 @@ def to_dtype_cvcuda(
 
 
 if CVCUDA_AVAILABLE:
-    _register_kernel_internal(to_dtype, cvcuda.Tensor)(to_dtype_cvcuda)
+    _to_dtype_cvcuda_registered = _register_kernel_internal(to_dtype, _import_cvcuda().Tensor)(_to_dtype_cvcuda)
 
 
 def sanitize_bounding_boxes(