Nigel
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

mindspore ops: solve segfault of PadAndShift when shift_idx is invalid

This commit is contained in:
liu lili 2023-02-10 11:30:15 +08:00
parent 6d49821c11
commit e677763f8b
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mindspore/ccsrc/plugin/device/cpu/kernel/pad_and_shift_cpu_kernel.cc
View File

@ -65,10 +65,10 @@ void PadAndShiftCpuKernelMod::LaunchKernel(const std::vector<AddressPtr> &inputs
T shift_idx = *reinterpret_cast<T *>(inputs[2]->addr);
T *output = reinterpret_cast<T *>(outputs[0]->addr);
if (shift_idx >= static_cast<T>(cum_sum_size_)) {
if (shift_idx < 0 || shift_idx >= static_cast<T>(cum_sum_size_)) {
MS_LOG(EXCEPTION) << "For '" << kernel_name_
<< "', shift index must be less than cumsum size, but got shift index: " << shift_idx
<< " and cumsum size: " << cum_sum_size_;
<< "', shift index must be large than 0 and less than cumsum size, but got shift index: "
<< shift_idx << " and cumsum size: " << cum_sum_size_;
}
size_t output_size = static_cast<size_t>(cum_sum_arr[cum_sum_size_ - 1]);
size_t shift_size = static_cast<size_t>(cum_sum_arr[shift_idx]);