From 011c48a8785a1ec28dd585313edc73c84e17ed45 Mon Sep 17 00:00:00 2001
From: wang_shaocong <wangshaocong1@huawei.com>
Date: Tue, 24 Nov 2020 14:28:37 +0800
Subject: [PATCH] fix bug of unsqueeze

---
 mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc | 1 +
 mindspore/lite/src/scheduler.cc                              | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc b/mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc
index 7a9110f467..06a02bd7a1 100644
--- 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
diff --git a/mindspore/lite/src/scheduler.cc b/mindspore/lite/src/scheduler.cc
index bc62e96ed0..d73d3c4e2c 100644
--- 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;
     }
   }