saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

GatherND - add Cuda support for int64 on opset 12 (#5531)

Browse source 
* support for int64

* per comments

Co-authored-by: Ethan Tao <ettao@microsoft.com>
This commit is contained in:
ytaous 2020-10-19 21:48:00 -07:00 committed by GitHub
parent 610676293e
commit 67968441e0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 8 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/core/providers/cuda/cuda_execution_provider.cc
View file

@ -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;
}

8
onnxruntime/core/providers/cuda/tensor/gather_nd.cc
View file

@ -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>);