Browse Source

fix bug of unsqueeze

tags/v1.1.0
wang_shaocong 5 years ago
parent
commit
011c48a878
2 changed files with 3 additions and 2 deletions
  1. +1
    -0
      mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc
  2. +2
    -2
      mindspore/lite/src/scheduler.cc

+ 1
- 0
mindspore/lite/src/runtime/kernel/arm/fp32/unsqueeze_fp32.cc View File

@@ -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

+ 2
- 2
mindspore/lite/src/scheduler.cc View File

@@ -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;
}
}


Loading…
Cancel
Save