diff --git a/test/prototype_transforms_kernel_infos.py b/test/prototype_transforms_kernel_infos.py index f5a2c6bb24c..52ff4565a55 100644 --- a/test/prototype_transforms_kernel_infos.py +++ b/test/prototype_transforms_kernel_infos.py @@ -66,6 +66,11 @@ def __init__( (("TestKernels", "test_against_reference"), torch.uint8, "cpu"): dict(atol=1e-5, rtol=0, agg_method="mean"), } +CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE = { + (("TestKernels", "test_cuda_vs_cpu"), dtype, "cuda"): dict(atol=atol, rtol=0) + for dtype, atol in [(torch.uint8, 1), (torch.float32, 1 / 255)] +} + def pil_reference_wrapper(pil_kernel): @functools.wraps(pil_kernel) @@ -319,7 +324,10 @@ def reference_inputs_resize_bounding_box(): sample_inputs_fn=sample_inputs_resize_image_tensor, reference_fn=reference_resize_image_tensor, reference_inputs_fn=reference_inputs_resize_image_tensor, - closeness_kwargs=DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + closeness_kwargs={ + **DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + **CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, + }, test_marks=[ xfail_jit_python_scalar_arg("size"), ], @@ -346,6 +354,7 @@ def reference_inputs_resize_bounding_box(): KernelInfo( F.resize_video, sample_inputs_fn=sample_inputs_resize_video, + closeness_kwargs=CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, ), ] ) @@ -1011,7 +1020,10 @@ def sample_inputs_resized_crop_video(): sample_inputs_fn=sample_inputs_resized_crop_image_tensor, reference_fn=reference_resized_crop_image_tensor, reference_inputs_fn=reference_inputs_resized_crop_image_tensor, - closeness_kwargs=DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + closeness_kwargs={ + **DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + **CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, + }, ), KernelInfo( F.resized_crop_bounding_box, @@ -1027,6 +1039,7 @@ def sample_inputs_resized_crop_video(): KernelInfo( F.resized_crop_video, sample_inputs_fn=sample_inputs_resized_crop_video, + closeness_kwargs=CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, ), ] ) @@ -1226,7 +1239,10 @@ def sample_inputs_perspective_video(): sample_inputs_fn=sample_inputs_perspective_image_tensor, reference_fn=pil_reference_wrapper(F.perspective_image_pil), reference_inputs_fn=reference_inputs_perspective_image_tensor, - closeness_kwargs=DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + closeness_kwargs={ + **DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + **CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, + }, ), KernelInfo( F.perspective_bounding_box, @@ -1242,6 +1258,7 @@ def sample_inputs_perspective_video(): KernelInfo( F.perspective_video, sample_inputs_fn=sample_inputs_perspective_video, + closeness_kwargs=CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, ), ] ) @@ -1442,7 +1459,10 @@ def sample_inputs_gaussian_blur_video(): KernelInfo( F.gaussian_blur_image_tensor, sample_inputs_fn=sample_inputs_gaussian_blur_image_tensor, - closeness_kwargs=DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + closeness_kwargs={ + **DEFAULT_PIL_REFERENCE_CLOSENESS_KWARGS, + **CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, + }, test_marks=[ xfail_jit_python_scalar_arg("kernel_size"), xfail_jit_python_scalar_arg("sigma"), @@ -1451,6 +1471,7 @@ def sample_inputs_gaussian_blur_video(): KernelInfo( F.gaussian_blur_video, sample_inputs_fn=sample_inputs_gaussian_blur_video, + closeness_kwargs=CUDA_VS_CPU_SINGLE_PIXEL_DIFFERENCE, ), ] )