From 317c2247aaaeddd81e8e4a0602dc269d59b85170 Mon Sep 17 00:00:00 2001 From: Difers <707065510@qq.com> Date: Fri, 25 Jul 2025 17:05:50 +0800 Subject: [PATCH] add PADDLE_ENFORCE_LE for unbind --- paddle/phi/kernels/impl/unbind_kernel_impl.h | 7 +++++++ paddle/phi/kernels/stride/unbind_kernel.cc | 7 +++++++ 2 files changed, 14 insertions(+) diff --git a/paddle/phi/kernels/impl/unbind_kernel_impl.h b/paddle/phi/kernels/impl/unbind_kernel_impl.h index 3e233a2038e480..bdb667f49731a3 100644 --- a/paddle/phi/kernels/impl/unbind_kernel_impl.h +++ b/paddle/phi/kernels/impl/unbind_kernel_impl.h @@ -24,6 +24,13 @@ void UnbindKernel(const Context& dev_ctx, const DenseTensor& x, int axis, std::vector outs) { + int64_t unbind_numel = x.numel(); + PADDLE_ENFORCE_LE(unbind_numel, + std::numeric_limits::max(), + ::common::errors::PreconditionNotMet( + "The number of proposals in unbind should be less than " + "2^31, but got %lld. Please check the input tensor. ", + unbind_numel)); auto x_dims = x.dims(); axis = axis < 0 ? x_dims.size() + axis : axis; diff --git a/paddle/phi/kernels/stride/unbind_kernel.cc b/paddle/phi/kernels/stride/unbind_kernel.cc index 0c7635aa53029c..1a2dcbb8e9f62d 100644 --- a/paddle/phi/kernels/stride/unbind_kernel.cc +++ b/paddle/phi/kernels/stride/unbind_kernel.cc @@ -32,6 +32,13 @@ void UnbindStridedKernel(const Context& dev_ctx, "FLAGS_use_stride_kernel is closed. Strided kernel " "be called, something wrong has happened!")); } + int64_t unbind_numel = x.numel(); + PADDLE_ENFORCE_LE(unbind_numel, + std::numeric_limits::max(), + ::common::errors::PreconditionNotMet( + "The number of proposals in unbind should be less than " + "2^31, but got %lld. Please check the input tensor. ", + unbind_numel)); int64_t num = static_cast(outs.size()); int64_t start = 0;