From 67968441e0e3e670326ea20af973afd3318d48dc Mon Sep 17 00:00:00 2001
From: ytaous <4484531+ytaous@users.noreply.github.com>
Date: Mon, 19 Oct 2020 21:48:00 -0700
Subject: [PATCH] GatherND - add Cuda support for int64 on opset 12 (#5531)

* support for int64

* per comments

Co-authored-by: Ethan Tao <ettao@microsoft.com>
---
 .../core/providers/cuda/cuda_execution_provider.cc        | 1 +
 onnxruntime/core/providers/cuda/tensor/gather_nd.cc       | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc
index 8c9d554984..fbc57899dc 100644
--- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc
+++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc
@@ -1840,6 +1840,7 @@ CUDAExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
 
     // none of the provided registries has a CUDA kernel for this node
     if (cuda_kernel_def == nullptr) {
+      LOGS_DEFAULT(WARNING) << "CUDA kernel not found in registries for Op type: " << node.OpType() << " node name: " << node.Name();
       continue;
     }
 
diff --git a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc
index fe2ac60d63..253c5d00a0 100644
--- a/onnxruntime/core/providers/cuda/tensor/gather_nd.cc
+++ b/onnxruntime/core/providers/cuda/tensor/gather_nd.cc
@@ -99,7 +99,13 @@ Status GatherNDBase::PrepareCompute(
       TIndex,                                                               \
       kCudaExecutionProvider,                                               \
       KernelDefBuilder()                                                    \
-          .TypeConstraint("T", DataTypeImpl::AllIEEEFloatTensorTypes())     \
+          .TypeConstraint("T",                                              \
+                          std::vector<MLDataType>{                          \
+                              DataTypeImpl::GetTensorType<float>(),         \
+                              DataTypeImpl::GetTensorType<double>(),        \
+                              DataTypeImpl::GetTensorType<MLFloat16>(),     \
+                              DataTypeImpl::GetTensorType<int64_t>(),       \
+                          })                                                \
           .TypeConstraint("Tind", DataTypeImpl::GetTensorType<TIndex>()),   \
       GatherND<TIndex>);