!2246 fix cpu argmax register

Merge pull request !2246 from kisnwang/fix-cpu-argmax
2020-06-18 15:29:52 +08:00 · 2020-06-18 15:29:52 +08:00 · 21ade66802
parent 74c3e15675 9d5502a4b6
commit 21ade66802
1 changed files with 1 additions and 1 deletions
--- 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