diff options
| author |  Vijay Vasudevan <vrv@google.com> | 2016-12-21 16:29:41 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2016-12-21 16:48:58 -0800 |
| commit | 21d8bb6dba6b2293901852c3fdbd113ee3805901 (patch) | |
| tree | dfe8cf2d33bbfd8a85a8efd893d90fa42051acc2 | |
| parent | 464a7e2315155fc71f98ea5e9ae8ef1bccb835e1 (diff) | |
Only register Reshape CPU kernels with type constraint of int32 on
shape parameter.
Change: 142713914
| -rw-r--r-- | tensorflow/core/kernels/reshape_op.cc | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/tensorflow/core/kernels/reshape_op.cc b/tensorflow/core/kernels/reshape_op.cc index 754eae7c46..6589a54624 100644 --- a/tensorflow/core/kernels/reshape_op.cc +++ b/tensorflow/core/kernels/reshape_op.cc @@ -18,7 +18,10 @@ limitations under the License. namespace tensorflow { -REGISTER_KERNEL_BUILDER(Name("Reshape").Device(DEVICE_CPU).HostMemory("shape"), +REGISTER_KERNEL_BUILDER(Name("Reshape") + .Device(DEVICE_CPU) + .HostMemory("shape") + .TypeConstraint<int32>("Tshape"), ReshapeOp); #define REGISTER_GPU_KERNEL(type) \ |