!20400 Fix the invalid input data for the ReduceXXX Op.

Merge pull request !20400 from liqiliang/master
This commit is contained in:
i-robot 2021-07-19 11:17:53 +00:00 committed by Gitee
commit a993f5a46a
1 changed files with 4 additions and 0 deletions

View File

@ -79,6 +79,10 @@ template <typename T>
bool ReduceCPUKernel<T>::Launch(const std::vector<kernel::AddressPtr> &inputs, const std::vector<kernel::AddressPtr> &,
const std::vector<kernel::AddressPtr> &outputs) {
size_t input_size = inputs[0]->size / sizeof(T);
if (input_size == 0) {
MS_LOG(EXCEPTION) << "Input data size is 0.";
}
auto input_addr = reinterpret_cast<T *>(inputs[0]->addr);
auto output_addr = reinterpret_cast<T *>(outputs[0]->addr);
if (axis_.empty() || input_shape_.empty() || input_shape_.size() == 1) {