diff --git a/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp b/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp index 99dc8a89293..46f59bba89f 100644 --- a/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp +++ b/kernels/portable/cpu/op_max_pool2d_with_indices_backward.cpp @@ -9,6 +9,8 @@ #include #include +#include + namespace torch { namespace executor { namespace native { @@ -171,6 +173,13 @@ Tensor& max_pool2d_with_indices_backward_out( static constexpr auto name = "max_pool2d_with_indices_backward.grad_input"; + // max_pool_backward_impl scatter-adds into grad_input (`grad_input_ptr[maxindex] += ...`), writing + // only the argmax positions. resize_tensor does not clear the buffer and the memory planner recycles + // arena allocations across ops and iterations, so every other element would accumulate onto stale + // data. ATen zeroes gradInput before dispatching the identical loop + // (aten/src/ATen/native/DilatedMaxPool2d.cpp). + memset(grad_input.mutable_data_ptr(), 0, grad_input.nbytes()); + ET_SWITCH_FLOATHBF16_TYPES(input.scalar_type(), ctx, name, CTYPE, [&]() { max_pool_backward_impl(grad_input, grad_output, indices); });