p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

fix cpu argmax register

This commit is contained in:
kswang 2020-06-17 20:09:39 +08:00
parent ee27718176
commit 9d5502a4b6
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h
View File

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