fix bug of unsqueeze

pull/8940/head
wang_shaocong 4 years ago
parent 5635e40b7c
commit 011c48a878

@ -100,4 +100,5 @@ kernel::LiteKernel *CpuUnsqueezeFp32KernelCreator(const std::vector<lite::Tensor
REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator) REG_KERNEL(kCPU, kNumberTypeFloat32, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator)
REG_KERNEL(kCPU, kNumberTypeInt32, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator) REG_KERNEL(kCPU, kNumberTypeInt32, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator)
REG_KERNEL(kCPU, kNumberTypeInt64, PrimitiveType_Unsqueeze, CpuUnsqueezeFp32KernelCreator)
} // namespace mindspore::kernel } // namespace mindspore::kernel

@ -355,8 +355,8 @@ kernel::SubGraphType Scheduler::GetKernelSubGraphType(const kernel::LiteKernel *
if (desc.data_type == kNumberTypeFloat16) { if (desc.data_type == kNumberTypeFloat16) {
return kernel::kCpuFP16SubGraph; return kernel::kCpuFP16SubGraph;
} else if (desc.data_type == kNumberTypeFloat32 || desc.data_type == kNumberTypeInt8 || } else if (desc.data_type == kNumberTypeFloat32 || desc.data_type == kNumberTypeInt8 ||
desc.data_type == kNumberTypeInt32 || desc.data_type == kNumberTypeBool || desc.data_type == kNumberTypeInt32 || desc.data_type == kNumberTypeInt64 ||
desc.data_type == kNumberTypeUInt8) { desc.data_type == kNumberTypeUInt8 || desc.data_type == kNumberTypeBool) {
return kernel::kCpuFP32SubGraph; return kernel::kCpuFP32SubGraph;
} }
} }

Loading…
Cancel
Save