Browse Source

fix cpu argmax register

tags/v0.5.0-beta
kswang 5 years ago
parent
commit
9d5502a4b6
1 changed files with 1 additions and 1 deletions
  1. +1
    -1
      mindspore/ccsrc/kernel/cpu/argmax_cpu_kernel.h

+ 1
- 1
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


Loading…
Cancel
Save