Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions paddle/phi/kernels/impl/unbind_kernel_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,13 @@ void UnbindKernel(const Context& dev_ctx,
const DenseTensor& x,
int axis,
std::vector<DenseTensor*> outs) {
int64_t unbind_numel = x.numel();
PADDLE_ENFORCE_LE(unbind_numel,
std::numeric_limits<int>::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;

Expand Down
7 changes: 7 additions & 0 deletions paddle/phi/kernels/stride/unbind_kernel.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>::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<int64_t>(outs.size());
int64_t start = 0;

Expand Down