simplify normalize testing into single test parameterize on input creation

Author: justincdavis

test/test_transforms_v2.py

@@ -5650,24 +5650,36 @@ def _reference_normalize_image(self, image, *, mean, std):
 
     @pytest.mark.parametrize(("mean", "std"), MEANS_STDS)
     @pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64])
+    @pytest.mark.parametrize(
+        "make_input",
+        [
+            make_image,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
+        ],
+    )
     @pytest.mark.parametrize("fn", [F.normalize, transform_cls_to_functional(transforms.Normalize)])
-    def test_correctness_image(self, mean, std, dtype, fn):
-        image = make_image(dtype=dtype)
+    def test_correctness_image(self, mean, std, dtype, make_input, fn):
+        if make_input == make_image_cvcuda and dtype != torch.float32:
+            pytest.skip("CVCUDA only supports float32 for normalize")
+
+        image = make_input(dtype=dtype)
 
         actual = fn(image, mean=mean, std=std)
-        expected = self._reference_normalize_image(image, mean=mean, std=std)
 
-        assert_equal(actual, expected)
+        if make_input == make_image_cvcuda:
+            image = F.cvcuda_to_tensor(image).to(device="cpu")
+            image = image.squeeze(0)
+            actual = F.cvcuda_to_tensor(actual).to(device="cpu")
+            actual = actual.squeeze(0)
 
-    @pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
-    @pytest.mark.parametrize(("mean", "std"), MEANS_STDS)
-    @pytest.mark.parametrize("fn", [F.normalize, transform_cls_to_functional(transforms.Normalize)])
-    def test_correctness_cvcuda(self, mean, std, fn):
-        image = make_image(batch_dims=(1,), dtype=torch.float32, device="cuda")
-        cvc_image = F.to_cvcuda_tensor(image)
-        actual = F._misc._normalize_cvcuda(cvc_image, mean=mean, std=std)
-        expected = fn(image, mean=mean, std=std)
-        torch.testing.assert_close(F.cvcuda_to_tensor(actual), expected, rtol=1e-7, atol=1e-7)
+        expected = self._reference_normalize_image(image, mean=mean, std=std)
+
+        if make_input == make_image_cvcuda:
+            torch.testing.assert_close(actual, expected, rtol=0, atol=1e-6)
+        else:
+            assert_equal(actual, expected)
 
 
 class TestClampBoundingBoxes:

torchvision/transforms/v2/_transform.py

@@ -8,7 +8,7 @@
 from torch import nn
 from torch.utils._pytree import tree_flatten, tree_unflatten
 from torchvision import tv_tensors
-from torchvision.transforms.v2._utils import check_type, has_any, is_pure_tensor
+from torchvision.transforms.v2._utils import check_type, has_any, is_cvcuda_tensor, is_pure_tensor
 from torchvision.utils import _log_api_usage_once
 
 from .functional._utils import _get_kernel
@@ -23,7 +23,7 @@ class Transform(nn.Module):
 
     # Class attribute defining transformed types. Other types are passed-through without any transformation
     # We support both Types and callables that are able to do further checks on the type of the input.
-    _transformed_types: tuple[type | Callable[[Any], bool], ...] = (torch.Tensor, PIL.Image.Image)
+    _transformed_types: tuple[type | Callable[[Any], bool], ...] = (torch.Tensor, PIL.Image.Image, is_cvcuda_tensor)
 
     def __init__(self) -> None:
         super().__init__()