add support for 1 channel float64

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -6807,15 +6807,25 @@ def test_invalid_dimensions(self):
         with pytest.raises(ValueError, match=r"pic should be 3 or 4 dimensional"):
             F.to_nvcv_tensor(torch.randint(0, 256, (1, 1, 3, 4, 4), dtype=torch.uint8, device="cuda"))
 
-    def test_unsupported_dtype_for_channels(self):
-        # Float64 is not supported
+    def test_float64_tensor_to_nvcv_tensor(self):
+        # Test single channel float64 (F64 format is supported)
+        img_data = torch.rand(1, 4, 4, dtype=torch.float64, device="cuda")
+        nvcv_img = F.to_nvcv_tensor(img_data)
+        assert nvcv_img is not None
+
+    def test_float64_rgb_not_supported(self):
+        # Test 3-channel float64 is NOT supported (no RGBf64 format in CV-CUDA)
         img_data = torch.rand(3, 4, 4, dtype=torch.float64, device="cuda")
         with pytest.raises(TypeError, match=r"Unsupported dtype"):
             F.to_nvcv_tensor(img_data)
 
     @pytest.mark.parametrize("num_channels", [1, 3])
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32, torch.float64])
     def test_round_trip(self, num_channels, dtype):
+        # Skip float64 for 3-channel (not supported by CV-CUDA)
+        if num_channels == 3 and dtype == torch.float64:
+            pytest.skip("float64 is not supported for 3-channel RGB images")
+
         # Setup: Create a tensor in CHW format (PyTorch standard)
         if dtype == torch.uint8:
             original_tensor = torch.randint(0, 256, (num_channels, 4, 4), dtype=dtype, device="cuda")
@@ -6832,9 +6842,13 @@ def test_round_trip(self, num_channels, dtype):
         assert torch.allclose(result_tensor, original_tensor, rtol=1e-5, atol=1e-7)
 
     @pytest.mark.parametrize("num_channels", [1, 3])
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32, torch.float64])
     @pytest.mark.parametrize("batch_size", [1, 2, 4])
     def test_round_trip_batched(self, num_channels, dtype, batch_size):
+        # Skip float64 for 3-channel (not supported by CV-CUDA)
+        if num_channels == 3 and dtype == torch.float64:
+            pytest.skip("float64 is not supported for 3-channel RGB images")
+
         # Setup: Create a batched tensor in NCHW format
         if dtype == torch.uint8:
             original_tensor = torch.randint(0, 256, (batch_size, num_channels, 4, 4), dtype=dtype, device="cuda")

torchvision/transforms/v2/functional/_type_conversion.py

@@ -61,6 +61,8 @@ def _infer_nvcv_format(img_tensor: torch.Tensor):
             return nvcv.Format.S32, img_tensor
         elif dtype == torch.float32:
             return nvcv.Format.F32, img_tensor
+        elif dtype == torch.float64:
+            return nvcv.Format.F64, img_tensor
         else:
             raise TypeError(f"Unsupported dtype {dtype} for single channel image")
 
@@ -71,6 +73,7 @@ def _infer_nvcv_format(img_tensor: torch.Tensor):
         elif dtype == torch.float32:
             return nvcv.Format.RGBf32, img_tensor
         else:
+            # 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.")