Modify the kernel declaration for Shrink op (#1554)

* Add capability for the input and output of Shrink op to share a commong buffer

* Cosmetic change
This commit is contained in:
Hariharan Seshadri 2019-08-05 13:21:04 -07:00 committed by GitHub
parent 6c271c63ac
commit ceb8f1c1a2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 1 deletions

View file

@ -11,7 +11,9 @@ namespace onnxruntime {
ONNX_CPU_OPERATOR_KERNEL(
Shrink,
9,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllNumericTensorTypes()),
KernelDefBuilder()
.MayInplace(0, 0)
.TypeConstraint("T", DataTypeImpl::AllNumericTensorTypes()),
Shrink);
namespace shrink_internal {

View file

@ -17,6 +17,7 @@ namespace cuda {
T, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.MayInplace(0, 0) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Shrink<T>);