fix return for _validate_cvcuda_dtype

Author: AntoineSimoulin

torchvision/transforms/v2/functional/_type_conversion.py

@@ -53,51 +53,32 @@ def to_image(inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray]) -> tv_tenso
 pil_to_tensor = _F.pil_to_tensor
 
 
-def _infer_cvcuda_format(img_tensor: torch.Tensor):
-    """Infer CV-CUDA format from tensor shape and dtype.
+def _validate_cvcuda_dtype(img_tensor: torch.Tensor) -> None:
+    """Validate that tensor dtype and channel count are supported by CV-CUDA.
 
     Args:
         img_tensor: Tensor with shape (H, W, C) where C is number of channels.
 
-    Returns:
-        tuple: (cvcuda_format, processed_tensor) where processed_tensor may have reduced dimensions
-               for single channel images.
-
     Raises:
         TypeError: If dtype is not supported for the given number of channels.
         ValueError: If number of channels is not 1 or 3.
     """
-    cvcuda = _import_cvcuda_modules()
-
     num_channels = img_tensor.shape[2]
     dtype = img_tensor.dtype
 
     # Handle single channel images
     if num_channels == 1:
-        if dtype == torch.uint8:
-            return cvcuda.Format.U8, img_tensor
-        elif dtype == torch.int16:
-            return cvcuda.Format.S16, img_tensor
-        elif dtype == torch.int32:
-            return cvcuda.Format.S32, img_tensor
-        elif dtype == torch.float32:
-            return cvcuda.Format.F32, img_tensor
-        elif dtype == torch.float64:
-            return cvcuda.Format.F64, img_tensor
-        else:
+        if dtype not in (torch.uint8, torch.int16, torch.int32, torch.float32, torch.float64):
             raise TypeError(f"Unsupported dtype {dtype} for single channel image")
 
     # Handle 3 channel images (defaults to RGB)
     elif num_channels == 3:
-        if dtype == torch.uint8:
-            return cvcuda.Format.RGB8, img_tensor
-        elif dtype == torch.float32:
-            return cvcuda.Format.RGBf32, img_tensor
-        else:
+        if dtype not in (torch.uint8, torch.float32):
             # Note: CV-CUDA does not support float64 for RGB images (only F64 for single-channel)
             raise TypeError(f"Unsupported dtype {dtype} for 3-channel image")
 
-    raise ValueError(f"Only 1 and 3 channel images are supported. Got {num_channels} channels.")
+    else:
+        raise ValueError(f"Only 1 and 3 channel images are supported. Got {num_channels} channels.")
 
 
 @torch.jit.unused
@@ -136,9 +117,9 @@ def to_cvcuda_tensor(pic) -> "cvcuda.Tensor":
     # Convert NCHW -> NHWC
     img_tensor = img_tensor.permute(0, 2, 3, 1)
 
-    # Infer format from the first image
+    # Validate dtype and channel count from the first image
     sample_img = img_tensor[0]
-    _infer_cvcuda_format(sample_img)
+    _validate_cvcuda_dtype(sample_img)
 
     # Convert to CV-CUDA tensor with NHWC layout
     return cvcuda.as_tensor(img_tensor.cuda().contiguous(), cvcuda.TensorLayout.NHWC)