!19537 fix allgather op select bug

Merge pull request !19537 from yuchaojie/r1.3_op_select
This commit is contained in:
i-robot 2021-07-07 06:14:04 +00:00 committed by Gitee
commit f03091b5a9
1 changed files with 1 additions and 1 deletions

View File

@ -45,7 +45,7 @@ std::string GetKernelFormat(const CNodePtr &kernel_node, size_t index) {
return format; return format;
} }
auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, index); auto input_shape = AnfAlgo::GetPrevNodeOutputInferShape(kernel_node, index);
if (op_name == kAllGatherOpName) { if (op_name == kAllGatherOpName && input_shape.size() <= kShape4dDims) {
auto pad_shape = trans::PaddingShapeTo4dDefault(input_shape); auto pad_shape = trans::PaddingShapeTo4dDefault(input_shape);
if (pad_shape[N_nchw] % kCubeSize != 0 || pad_shape[C_nchw] % kCubeSize != 0) { if (pad_shape[N_nchw] % kCubeSize != 0 || pad_shape[C_nchw] % kCubeSize != 0) {
return kOpFormat_DEFAULT; return kOpFormat_DEFAULT;