From 9d5502a4b697bcdd2d23918901103f89d9715918 Mon Sep 17 00:00:00 2001 From: kswang Date: Wed, 17 Jun 2020 20:09:39 +0800 Subject: [PATCH] fix cpu argmax register --- mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h b/mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h index 16344d68179..aae7435c5c6 100644 --- a/mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h +++ b/mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h @@ -37,7 +37,7 @@ class ArgmaxCPUKernel : public CPUKernel { size_t batch_size_{0}; }; -MS_REG_CPU_KERNEL(Argmax, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32), +MS_REG_CPU_KERNEL(Argmax, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeInt32), ArgmaxCPUKernel); } // namespace kernel } // namespace mindspore