1From 19edf6b71a81c21bcaa15e6e564fc6ad56a188da Mon Sep 17 00:00:00 2001
2From: rabbit-fgh <2955722401@qq.com>
3Date: Thu, 22 Aug 2024 11:22:18 +0800
4Subject: [PATCH] fix split problems
5
6---
7 .../plugin/device/cpu/kernel/nnacl/infer/split_infer.c | 10 ++++++----
8 1 file changed, 6 insertions(+), 4 deletions(-)
9
10diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/infer/split_infer.c b/mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/infer/split_infer.c
11index b017e259..1d3aabe1 100644
12--- a/mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/infer/split_infer.c
13+++ b/mindspore/ccsrc/plugin/device/cpu/kernel/nnacl/infer/split_infer.c
14@@ -34,11 +34,13 @@ int UpdateSplitSize(const TensorC *const *inputs, size_t inputs_size, SplitParam
15   }
16   if (param->split_count_ == 0) {
17     const TensorC *input = inputs[0];
18-    if (input->shape_[param->split_dim_] % param->num_split_ != 0) {
19-      return NNACL_ERR;
20-    }
21+    int32_t split_chunk_size = UP_DIV(input->shape_[param->split_dim_], param->num_split_);
22     for (int i = 0; i < param->num_split_; ++i) {
23-      param->split_sizes_[i] = input->shape_[param->split_dim_] / param->num_split_;
24+      if (i != param->num_split_ - 1) {
25+        param->split_sizes_[i] = split_chunk_size;
26+      } else {
27+        param->split_sizes_[i] = input->shape_[param->split_dim_] - split_chunk_size * i;
28+      }
29     }
30   }
31   return NNACL_OK;
32-- 
332.45.1.windows.1
34
35