From aefec228628ddf613d2006c055c57c7aad004adb Mon Sep 17 00:00:00 2001 From: TeslaZhao Date: Mon, 23 Aug 2021 14:08:27 +0800 Subject: [PATCH] Fix a bug of strided_slice op, about the axes parameter access memory out of bounds (#35062) --- paddle/fluid/operators/strided_slice_op.cc | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/paddle/fluid/operators/strided_slice_op.cc b/paddle/fluid/operators/strided_slice_op.cc index d53ab914db4..a1b5ca0f6a6 100644 --- a/paddle/fluid/operators/strided_slice_op.cc +++ b/paddle/fluid/operators/strided_slice_op.cc @@ -62,6 +62,20 @@ class StridedSliceOp : public framework::OperatorWithKernel { auto ends_size = ends.size(); auto strides_size = strides.size(); + for (size_t i = 0; i < axes.size(); ++i) { + PADDLE_ENFORCE_GE(axes[i], 0, + platform::errors::InvalidArgument( + "The axis should be greater than or equal to 0." + "But received %d of axes[%d]", + axes[i], i)); + PADDLE_ENFORCE_LT( + axes[i], in_dims.size(), + platform::errors::InvalidArgument( + "The axes should be less than or equal to input tensor's rank." + "But received %d of axes[%d], input tensor shape [%d]", + axes[i], i, in_dims.size())); + } + if (ctx->HasInputs("StartsTensorList")) { auto StartsTensorList = ctx->Inputs("StartsTensorList"); PADDLE_ENFORCE_GT( -- GitLab