diff --git a/src/ATen/native/xpu/sycl/GridSampler.cpp b/src/ATen/native/xpu/sycl/GridSampler.cpp index 4e4c094b3f9..b31154c85bf 100644 --- a/src/ATen/native/xpu/sycl/GridSampler.cpp +++ b/src/ATen/native/xpu/sycl/GridSampler.cpp @@ -682,8 +682,7 @@ void grid_sampler_2d_backward_kernel( int64_t padding_mode, bool align_corners, std::array output_mask) { - check_grid_sampler_common(input, grid); - check_grid_sampler_2d(input, grid); + check_grid_sampler_2d_backward(input, grid, grad_output); globalContext().alertNotDeterministic("grid_sampler_2d_backward_xpu"); auto N = input.size(0); @@ -1491,8 +1490,13 @@ void grid_sampler_3d_backward_kernel( std::array output_mask) { // See NOTE [ grid_sampler Native Functions ]. // Add checks here in case this is called instead of grid_sampler. - check_grid_sampler_common(input, grid); - check_grid_sampler_3d(input, grid, interpolation_mode); + check_grid_sampler_3d_backward(input, grid, grad_output); + TORCH_CHECK_NOT_IMPLEMENTED( + static_cast(interpolation_mode) != + GridSamplerInterpolation::Bicubic, + "grid_sampler(): bicubic interpolation with 5D input is not " + "implemented for ", + input.device().type()); globalContext().alertNotDeterministic("grid_sampler_3d_backward_xpu"); auto input_requires_grad = output_mask[0]; @@ -1506,7 +1510,7 @@ void grid_sampler_3d_backward_kernel( at::ScalarType::BFloat16, at::ScalarType::Half, input.scalar_type(), - "grid_sampler_2d_backward_xpu", + "grid_sampler_3d_backward_xpu", [&] { if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) && canUse32BitIndexMath(grad_output)) { diff --git a/test/xpu/test_nn_xpu.py b/test/xpu/test_nn_xpu.py index 80491dbb1af..35ac0066dff 100644 --- a/test/xpu/test_nn_xpu.py +++ b/test/xpu/test_nn_xpu.py @@ -14201,6 +14201,37 @@ def test_save_lstm_compatibility(self, device): result = loaded_model(x) self.assertEqual(result, expected) + def test_grid_sample_backward_error_checking(self, device): + input = torch.empty(1, 1, 2, 2, device=device) + grid = torch.empty(1, 1, 1, 2, device=device) + grad_output = torch.empty(1, 1, 1, 1, device=device) + + # assert no error + torch.ops.aten.grid_sampler_2d_backward( + grad_output, input, grid, 0, 0, False, (True, True) + ) + + with self.assertRaisesRegex(ValueError, "expected grad_output to have sizes"): + invalid_grad_output = torch.empty(2, 1, 1, 1, device=device) + torch.ops.aten.grid_sampler_2d_backward( + invalid_grad_output, input, grid, 0, 0, False, (True, True) + ) + + input = torch.empty(1, 1, 2, 2, 2, device=device) + grid = torch.empty(1, 1, 1, 1, 3, device=device) + grad_output = torch.empty(1, 1, 1, 1, 1, device=device) + + # assert no error + torch.ops.aten.grid_sampler_3d_backward( + grad_output, input, grid, 0, 0, False, (True, True) + ) + + with self.assertRaisesRegex(ValueError, "expected grad_output to have sizes"): + invalid_grad_output = torch.empty(2, 1, 1, 1, 1, device=device) + torch.ops.aten.grid_sampler_3d_backward( + invalid_grad_output, input, grid, 0, 0, False, (True, True) + ) + @onlyOn(["cuda", "xpu"]) @tf32_on_and_off(0.005) def test_grid_sample_large(self, device):