update invert with new PR comments revisions

Author: justincdavis

test/test_transforms_v2.py

@@ -5874,23 +5874,26 @@ def test_functional_signature(self, kernel, input_type):
     def test_transform(self, make_input):
         check_transform(transforms.RandomInvert(p=1), make_input())
 
+    @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.invert, transform_cls_to_functional(transforms.RandomInvert, p=1)])
-    def test_correctness_image(self, fn):
-        image = make_image(dtype=torch.uint8, device="cpu")
+    def test_correctness_image(self, make_input, fn):
+        image = make_input(dtype=torch.uint8, device="cpu")
 
         actual = fn(image)
-        expected = F.to_image(F.invert(F.to_pil_image(image)))
 
-        assert_equal(actual, expected)
+        if make_input is make_image_cvcuda:
+            image = cvcuda_to_pil_compatible_tensor(image)
+
+        expected = F.to_image(F.invert(F.to_pil_image(image)))
 
-    @pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
-    @pytest.mark.parametrize("fn", [F.invert, transform_cls_to_functional(transforms.RandomInvert, p=1)])
-    def test_correctness_cvcuda(self, dtype, fn):
-        image = make_image(batch_dims=(1,), dtype=dtype, device="cuda")
-        cv_image = F.to_cvcuda_tensor(image)
-        actual = F.cvcuda_to_tensor(fn(cv_image))
-        expected = F.invert_image(image)
         assert_equal(actual, expected)
 
 

torchvision/transforms/v2/functional/_color.py

@@ -690,13 +690,15 @@ def invert_video(video: torch.Tensor) -> torch.Tensor:
     return invert_image(video)
 
 
-if _CVCUDA_AVAILABLE:
+if CVCUDA_AVAILABLE:
     _invert_cvcuda_tensors = {}
 
 
 def _invert_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
     cvcuda = _import_cvcuda()
 
+    # save the tensors into a dictionary only if CV-CUDA is actually used
+    # we save these here, since they are static and small in size
     if "base" not in _invert_cvcuda_tensors:
         _invert_cvcuda_tensors["base"] = cvcuda.as_tensor(
             torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device="cuda").reshape(1, 1, 1, 3).contiguous(), "NHWC"
@@ -722,7 +724,7 @@ def _invert_cvcuda(image: "cvcuda.Tensor") -> "cvcuda.Tensor":
     return cvcuda.normalize(image, base=base, scale=scale, globalscale=1.0, globalshift=shift)
 
 
-if _CVCUDA_AVAILABLE:
+if CVCUDA_AVAILABLE:
     _register_kernel_internal(invert, _import_cvcuda().Tensor)(_invert_cvcuda)