saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-25 19:48:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix registration for GatherGrad (#5382)

Browse source 
* Fix registration for GatherGrad to fix GatherGradOpTest.GatherGrad_axis0_indices2d_half.

* Fix GatherGrad registration for CUDA also.
This commit is contained in:
Sergii Dymchenko 2020-10-09 11:57:50 -07:00 committed by GitHub
parent 1cceefc7d4
commit 3a9a1a4ef1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 5 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -446,7 +446,7 @@ void RegisterTrainingOpSchemas() {
"Constrain input shape to integer tensors.")
.TypeConstraint(
"T",
OpSchema::all_tensor_types_with_bfloat(),
{"tensor(float16)", "tensor(float)", "tensor(double)"},
"Constrain input and output types to float tensors.")
.TypeConstraint(
"Tind",

3
orttraining/orttraining/training_ops/cpu/tensor/gather_grad.cc
View file

@ -15,7 +15,8 @@ ONNX_OPERATOR_KERNEL_EX(
kCpuExecutionProvider,
KernelDefBuilder()
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<double>()})
.TypeConstraint("Tind", std::vector<MLDataType>{
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),

3
orttraining/orttraining/training_ops/cuda/tensor/gather_grad.cc
View file

@ -16,7 +16,8 @@ ONNX_OPERATOR_KERNEL_EX(
KernelDefBuilder()
.InputMemoryType<OrtMemTypeCPUInput>(0)
.TypeConstraint("I", DataTypeImpl::GetTensorType<int64_t>())
.TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
.TypeConstraint("T", {DataTypeImpl::GetTensorType<float>(),
DataTypeImpl::GetTensorType<MLFloat16>()})
.TypeConstraint("Tind", std::vector<MLDataType>{
DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),