From 32ad2aea13f05c2d77bacb50b83cf4df13273de4 Mon Sep 17 00:00:00 2001
From: zong-shuai <zongshuai1@huawei.com>
Date: Sat, 29 Jan 2022 15:53:53 +0800
Subject: [PATCH] debug_topk

---
 .../kernel_compiler/gpu/arrays/topk_gpu_kernel.cc | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc
index 3066c89ebd1..0566e7270ba 100644
--- a/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc
+++ b/mindspore/ccsrc/backend/kernel_compiler/gpu/arrays/topk_gpu_kernel.cc
@@ -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