strenghen condition for round trip test

Author: AntoineSimoulin

test/test_transforms_v2.py

@@ -6801,13 +6801,15 @@ def test_round_trip(self, dtype):
             original_tensor = torch.rand(3, 4, 4, dtype=dtype, device="cuda")
 
         # Execute: Convert to NVCV and back to tensor
-        # CHW -> (to_nvcv_tensor) -> NVCV NHWC -> (nvcv_to_tensor) -> CHW
+        # CHW -> (to_nvcv_tensor) -> NVCV NHWC -> (nvcv_to_tensor) -> NCHW
         nvcv_tensor = F.to_nvcv_tensor(original_tensor)
         result_tensor = F.nvcv_to_tensor(nvcv_tensor)
 
-        # Assert: The round-trip conversion preserves the original tensor
-        # Use allclose for robust comparison that handles floating-point precision
-        assert torch.allclose(result_tensor, original_tensor, rtol=1e-5, atol=1e-7)
+        # Remove batch dimension that was added during conversion since original was unbatched
+        result_tensor = result_tensor.squeeze(0)
+
+        # Assert: The round-trip conversion preserves the original tensor exactly
+        torch.testing.assert_close(result_tensor, original_tensor, rtol=0, atol=0)
 
     @pytest.mark.parametrize("dtype", [torch.uint8, torch.float32])
     @pytest.mark.parametrize("batch_size", [1, 2, 4])
@@ -6823,9 +6825,8 @@ def test_round_trip_batched(self, dtype, batch_size):
         nvcv_tensor = F.to_nvcv_tensor(original_tensor)
         result_tensor = F.nvcv_to_tensor(nvcv_tensor)
 
-        # Assert: The round-trip conversion preserves the original batched tensor
-        # Use allclose for robust comparison that handles floating-point precision
-        assert torch.allclose(result_tensor, original_tensor, rtol=1e-5, atol=1e-7)
+        # Assert: The round-trip conversion preserves the original batched tensor exactly
+        torch.testing.assert_close(result_tensor, original_tensor, rtol=0, atol=0)
         # Also verify batch size is preserved
         assert result_tensor.shape[0] == batch_size