add tv_tensors.Image in tests

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -6746,6 +6746,7 @@ def test_functional_error(self):
 @pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
 @needs_cuda
 class TestToCVCUDATensor:
+    @pytest.mark.parametrize("image_type", (torch.Tensor, tv_tensors.Image))
     @pytest.mark.parametrize("dtype", [torch.uint8, torch.uint16, torch.float32, torch.float64])
     @pytest.mark.parametrize("device", cpu_and_cuda())
     @pytest.mark.parametrize("color_space", ["RGB", "GRAY"])
@@ -6754,30 +6755,24 @@ class TestToCVCUDATensor:
         "fn",
         [F.to_cvcuda_tensor, transform_cls_to_functional(transforms.ToCVCUDATensor)],
     )
-    def test_functional_and_transform(self, dtype, device, color_space, batch_dims, fn):
-        input = make_image_tensor(dtype=dtype, device=device, color_space=color_space, batch_dims=batch_dims)
-        output = fn(input)
+    def test_functional_and_transform(self, image_type, dtype, device, color_space, batch_dims, fn):
+        image = make_image(dtype=dtype, device=device, color_space=color_space, batch_dims=batch_dims)
+        if image_type is torch.Tensor:
+            image = image.as_subclass(torch.Tensor)
+            assert is_pure_tensor(image)
+        output = fn(image)
 
         assert isinstance(output, cvcuda.Tensor)
-        assert F.get_size(output) == F.get_size(input)
+        assert F.get_size(output) == F.get_size(image)
         assert output is not None
 
     def test_invalid_input_type(self):
-        with pytest.raises(TypeError, match=r"inpt should be `torch.Tensor`"):
+        with pytest.raises(TypeError, match=r"inpt should be ``torch.Tensor``"):
             F.to_cvcuda_tensor("invalid_input")
 
     def test_invalid_dimensions(self):
         with pytest.raises(ValueError, match=r"pic should be 4 dimensional"):
-            img_data = torch.randint(
-                0,
-                256,
-                (
-                    3,
-                    12,
-                    34,
-                ),
-                dtype=torch.uint8,
-            )
+            img_data = torch.randint(0, 256, (3, 1, 3), dtype=torch.uint8)
             img_data = img_data.cuda()
             F.to_cvcuda_tensor(img_data)
 

torchvision/transforms/v2/functional/_type_conversion.py

@@ -39,8 +39,8 @@ def to_cvcuda_tensor(inpt: torch.Tensor) -> "cvcuda.Tensor":
     cvcuda = _import_cvcuda()
     if not torch.jit.is_scripting() and not torch.jit.is_tracing():
         _log_api_usage_once(to_cvcuda_tensor)
-    if not isinstance(inpt, torch.Tensor):
-        raise TypeError(f"inpt should be ``torch.Tensor``. Got {type(inpt)}.")
+    if not isinstance(inpt, (torch.Tensor, tv_tensors.Image)):
+        raise TypeError(f"inpt should be ``torch.Tensor`` or ``tv_tensors.Image``. Got {type(inpt)}.")
     if inpt.ndim != 4:
         raise ValueError(f"pic should be 4 dimensional. Got {inpt.ndim} dimensions.")
     # Convert to NHWC as CVCUDA transforms do not support NCHW