add int16/int32 dtypes in round trip tests

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -6845,15 +6845,15 @@ def test_float64_rgb_not_supported(self):
             F.to_cvcuda_tensor(img_data)
 
     @pytest.mark.parametrize("num_channels", [1, 3])
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32, torch.float64])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.int32, 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")
+        # Skip int16/int32/float64 for 3-channel (only supported for single-channel)
+        if num_channels == 3 and dtype in (torch.int16, torch.int32, torch.float64):
+            pytest.skip(f"{dtype} is only supported for single-channel images")
 
         # Setup: Create a tensor in CHW format (PyTorch standard)
         # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        if dtype == torch.uint8:
+        if dtype in (torch.uint8, torch.int16, torch.int32):
             original_tensor = torch.randint(0, 256, (num_channels, 4, 4), dtype=dtype)
         else:
             original_tensor = torch.rand(num_channels, 4, 4, dtype=dtype)
@@ -6871,16 +6871,16 @@ def test_round_trip(self, num_channels, dtype):
         torch.testing.assert_close(result_tensor, original_tensor, rtol=0, atol=0)
 
     @pytest.mark.parametrize("num_channels", [1, 3])
-    @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32, torch.float64])
+    @pytest.mark.parametrize("dtype", [torch.uint8, torch.int16, torch.int32, 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")
+        # Skip int16/int32/float64 for 3-channel (only supported for single-channel)
+        if num_channels == 3 and dtype in (torch.int16, torch.int32, torch.float64):
+            pytest.skip(f"{dtype} is only supported for single-channel images")
 
         # Setup: Create a batched tensor in NCHW format
         # Create tensor on CPU first, then move to CUDA to avoid CUDA context issues
-        if dtype == torch.uint8:
+        if dtype in (torch.uint8, torch.int16, torch.int32):
             original_tensor = torch.randint(0, 256, (batch_size, num_channels, 4, 4), dtype=dtype)
         else:
             original_tensor = torch.rand(batch_size, num_channels, 4, 4, dtype=dtype)