fix 3 dimensions check

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -6768,17 +6768,31 @@ def test_invalid_input_type(self):
             F.to_cvcuda_tensor("invalid_input")
 
     def test_invalid_dimensions(self):
-        with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
+        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 = img_data.cuda()
+            F.to_cvcuda_tensor(img_data)
+
+        with pytest.raises(ValueError, match=r"pic should be 4 dimensional"):
             img_data = torch.randint(0, 256, (4,), dtype=torch.uint8)
             img_data = img_data.cuda()
             F.to_cvcuda_tensor(img_data)
 
-        with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
+        with pytest.raises(ValueError, match=r"pic should be 4 dimensional"):
             img_data = torch.randint(0, 256, (4, 4), dtype=torch.uint8)
             img_data = img_data.cuda()
             F.to_cvcuda_tensor(img_data)
 
-        with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
+        with pytest.raises(ValueError, match=r"pic should be 4 dimensional"):
             img_data = torch.randint(0, 256, (1, 1, 3, 4, 4), dtype=torch.uint8)
             img_data = img_data.cuda()
             F.to_cvcuda_tensor(img_data)

torchvision/transforms/v2/functional/_type_conversion.py

@@ -41,8 +41,8 @@ def to_cvcuda_tensor(inpt: torch.Tensor) -> "cvcuda.Tensor":
         _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 inpt.ndim not in [3, 4]:
-        raise ValueError(f"pic should be 3 or 4 dimensional. Got {inpt.ndim} dimensions.")
+    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
     inpt = inpt.permute(0, 2, 3, 1)
     return cvcuda.as_tensor(inpt.cuda().contiguous(), cvcuda.TensorLayout.NHWC)
@@ -57,8 +57,8 @@ def cvcuda_to_tensor(cvcuda_img: "cvcuda.Tensor") -> torch.Tensor:
     if not isinstance(cvcuda_img, cvcuda.Tensor):
         raise TypeError(f"cvcuda_img should be `cvcuda.Tensor`. Got {type(cvcuda_img)}.")
     cuda_tensor = torch.as_tensor(cvcuda_img.cuda(), device="cuda")
-    if cvcuda_img.ndim not in [3, 4]:
-        raise ValueError(f"Image should be 3 or 4 dimensional. Got {cuda_tensor.ndim} dimensions.")
+    if cvcuda_img.ndim != 4:
+        raise ValueError(f"Image should be 4 dimensional. Got {cuda_tensor.ndim} dimensions.")
     # Convert to NCHW layout from CVCUDA default NHWC
     img = cuda_tensor.permute(0, 3, 1, 2)
     return img