fix bug of unsqueeze

4 years ago · 011c48a878
parent 5635e40b7c
commit 011c48a878
2 changed files with 3 additions and 2 deletions
--- a/mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc
+++ b/mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc
@ -100,4 +100,5 @@ kernel::LiteKernel *CpuUnsqueezeFp32KernelCreator(const std::vector<lite::Tensor

 REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator)
 REG_KERNEL(kCPU, kNumberTypeInt32, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator)
+REG_KERNEL(kCPU, kNumberTypeInt64, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator)
 }  // namespace mindspore::kernel
--- a/mindspore/lite/src/scheduler.cc
+++ b/mindspore/lite/src/scheduler.cc
@ -355,8 +355,8 @@ kernel::SubGraphType Scheduler::GetKernelSubGraphType(const kernel::LiteKernel *
    if (desc.data_type == kNumberTypeFloat16) {
      return kernel::kCpuFP16SubGraph;
    } else if (desc.data_type == kNumberTypeFloat32 || desc.data_type == kNumberTypeInt8 ||
-               desc.data_type == kNumberTypeInt32 || desc.data_type == kNumberTypeBool ||
-               desc.data_type == kNumberTypeUInt8) {
+               desc.data_type == kNumberTypeInt32 || desc.data_type == kNumberTypeInt64 ||
+               desc.data_type == kNumberTypeUInt8 || desc.data_type == kNumberTypeBool) {
      return kernel::kCpuFP32SubGraph;
    }
  }