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

!43608 fix gather with big shape 

Merge pull request !43608 from 范吉斌/fix_gather
This commit is contained in:
i-robot 2022-10-11 02:21:09 +00:00 committed by Gitee
parent efc0839ef2 4db51463f0
commit 2d7af0e080
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mindspore/ccsrc/plugin/device/cpu/kernel/gather_cpu_kernel.h
View File

@ -64,8 +64,9 @@ class GatherCpuKernelMod : public NativeCpuKernelMod {
if (is_dynamic_shape_) {
input_size_list_.push_back(axis_type_size_);
}
auto output_size = std::accumulate(output_shape_.begin(), output_shape_.end(), 1, std::multiplies{});
output_size_list_.push_back(LongToSize(output_size) * input_type_size_);
auto output_size =
std::accumulate(output_shape_.begin(), output_shape_.end(), static_cast<size_t>(1), std::multiplies{});
output_size_list_.push_back(output_size * input_type_size_);
}
private: