mindspore-Ecosystem
/
mindspore
21
0
Fork 27
Code Issues Pull Requests 6 Projects Releases Wiki Activity

debug_topk

This commit is contained in:
zong-shuai 2022-01-29 15:53:53 +08:00
parent c72a1146fb
commit 32ad2aea13
1 changed files with 15 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc
View File

@ -33,5 +33,20 @@ MS_REG_GPU_KERNEL_TWO(TopK,
.AddOutputAttr(kNumberTypeFloat16)
.AddOutputAttr(kNumberTypeInt32),
TopKGpuKernel, half, int)
MS_REG_GPU_KERNEL_TWO(TopK,
KernelAttr()
.AddInputAttr(kNumberTypeFloat32)
.AddInputAttr(kNumberTypeInt64)
.AddOutputAttr(kNumberTypeFloat32)
.AddOutputAttr(kNumberTypeInt32),
TopKGpuKernel, float, int)
MS_REG_GPU_KERNEL_TWO(TopK,
KernelAttr()
.AddInputAttr(kNumberTypeFloat16)
.AddInputAttr(kNumberTypeInt64)
.AddOutputAttr(kNumberTypeFloat16)
.AddOutputAttr(kNumberTypeInt32),
TopKGpuKernel, half, int)
} // namespace kernel
} // namespace mindspore