Fix: update center crop

Author: justincdavis

test/test_transforms_v2.py

@@ -4965,6 +4965,9 @@ def test_kernel_video(self):
             make_segmentation_mask,
             make_video,
             make_keypoints,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
         ],
     )
     def test_functional(self, make_input):
@@ -4980,6 +4983,11 @@ def test_functional(self, make_input):
             (F.center_crop_mask, tv_tensors.Mask),
             (F.center_crop_video, tv_tensors.Video),
             (F.center_crop_keypoints, tv_tensors.KeyPoints),
+            pytest.param(
+                F._geometry._center_crop_cvcuda,
+                _import_cvcuda().Tensor,
+                marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA"),
+            ),
         ],
     )
     def test_functional_signature(self, kernel, input_type):
@@ -4995,6 +5003,9 @@ def test_functional_signature(self, kernel, input_type):
             make_segmentation_mask,
             make_video,
             make_keypoints,
+            pytest.param(
+                make_image_cvcuda, marks=pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+            ),
         ],
     )
     def test_transform(self, make_input):
@@ -5010,6 +5021,17 @@ def test_image_correctness(self, output_size, fn):
 
         assert_equal(actual, expected)
 
+    @pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="test requires CVCUDA")
+    @pytest.mark.parametrize("output_size", OUTPUT_SIZES)
+    @pytest.mark.parametrize("fn", [F.center_crop, transform_cls_to_functional(transforms.CenterCrop)])
+    def test_cvcuda_correctness(self, output_size, fn):
+        image = make_image_cvcuda(self.INPUT_SIZE, dtype=torch.uint8, device="cuda")
+
+        actual = fn(image, output_size)
+        expected = F.center_crop(F.cvcuda_to_tensor(image), output_size)
+
+        assert_equal(F.cvcuda_to_tensor(actual), expected)
+
     def _reference_center_crop_bounding_boxes(self, bounding_boxes, output_size):
         image_height, image_width = bounding_boxes.canvas_size
         if isinstance(output_size, int):
@@ -5081,33 +5103,6 @@ def test_keypoints_correctness(self, output_size, dtype, device, fn):
         assert_equal(actual, expected)
 
 
-@pytest.mark.skipif(not CVCUDA_AVAILABLE, reason="cvcuda not available")
-@needs_cuda
-class TestCenterCropCVCUDA:
-    def test_functional(self):
-        check_functional(
-            F.center_crop,
-            make_image_cvcuda(TestCenterCrop.INPUT_SIZE, batch_dims=(1,)),
-            output_size=TestCenterCrop.OUTPUT_SIZES[0],
-        )
-
-    def test_functional_signature(self):
-        check_functional_kernel_signature_match(F.center_crop, kernel=F.center_crop_cvcuda, input_type=cvcuda.Tensor)
-
-    @pytest.mark.parametrize("output_size", TestCenterCrop.OUTPUT_SIZES)
-    def test_functional_correctness(self, output_size):
-        image = make_image_cvcuda(TestCenterCrop.INPUT_SIZE, batch_dims=(1,))
-        actual = F.center_crop(image, output_size)
-        expected = F.center_crop(F.cvcuda_to_tensor(image), output_size)
-        assert_equal(F.cvcuda_to_tensor(actual), expected)
-
-    def test_transform(self):
-        check_transform(
-            transforms.CenterCrop(TestCenterCrop.OUTPUT_SIZES[0]),
-            make_image_cvcuda(TestCenterCrop.INPUT_SIZE, batch_dims=(1,)),
-        )
-
-
 class TestPerspective:
     COEFFICIENTS = [
         [1.2405, 0.1772, -6.9113, 0.0463, 1.251, -5.235, 0.00013, 0.0018],

torchvision/transforms/v2/functional/_geometry.py

@@ -2718,7 +2718,7 @@ def center_crop_video(video: torch.Tensor, output_size: list[int]) -> torch.Tens
     return center_crop_image(video, output_size)
 
 
-def center_crop_cvcuda(
+def _center_crop_cvcuda(
     image: "cvcuda.Tensor",
     output_size: list[int],
 ) -> "cvcuda.Tensor":
@@ -2754,7 +2754,9 @@ def center_crop_cvcuda(
 
 
 if CVCUDA_AVAILABLE:
-    _register_kernel_internal(center_crop, cvcuda.Tensor)(center_crop_cvcuda)
+    _center_crop_cvcuda_registered = _register_kernel_internal(center_crop, _import_cvcuda().Tensor)(
+        _center_crop_cvcuda
+    )
 
 
 def resized_crop(