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

!29284 add int64 to k_type 

Merge pull request !29284 from zong_shuai/topk_debug
This commit is contained in:
i-robot 2022-01-19 06:36:38 +00:00 committed by Gitee
parent 754b92dd00 e123504ef7
commit b3dde2fe0e
No known key found for this signature in database
GPG Key ID: 173E9B9CA92EEF8F
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