diff --git a/runtime/core/exec_aten/util/tensor_util.h b/runtime/core/exec_aten/util/tensor_util.h index 9b490da244c..26b97e5a7a2 100644 --- a/runtime/core/exec_aten/util/tensor_util.h +++ b/runtime/core/exec_aten/util/tensor_util.h @@ -1212,6 +1212,12 @@ ET_NODISCARD inline Error resize_tensor( std::array new_sizes_casted{}; size_t new_sizes_ndim = new_sizes.size(); + ET_CHECK_OR_RETURN_ERROR( + new_sizes_ndim <= kTensorDimensionLimit, + InvalidArgument, + "new_sizes_ndim %zu is greater than kTensorDimensionLimit %zu", + new_sizes_ndim, + kTensorDimensionLimit); for (const auto i : c10::irange(new_sizes_ndim)) { new_sizes_casted[i] = static_cast(new_sizes[i]);