aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/ops/nn_ops.cc
diff options
context:
space:
mode:
authorGravatar Jianwei Xie <xiejw@google.com>2018-01-24 10:02:35 -0800
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-01-24 10:06:06 -0800
commitd9f93c42a50b1f1401d9c186eac0ae8dc9093c3b (patch)
tree178d1a692f56580c266139642b5a1d0d155c477e /tensorflow/core/ops/nn_ops.cc
parent7b62a71e2d46c148df7d5704972f4592bc5e0f1b (diff)
Merge changes from github.
PiperOrigin-RevId: 183100142
Diffstat (limited to 'tensorflow/core/ops/nn_ops.cc')
-rw-r--r--tensorflow/core/ops/nn_ops.cc8
1 files changed, 8 insertions, 0 deletions
diff --git a/tensorflow/core/ops/nn_ops.cc b/tensorflow/core/ops/nn_ops.cc
index 536fc7c0c1..3f72b41569 100644
--- a/tensorflow/core/ops/nn_ops.cc
+++ b/tensorflow/core/ops/nn_ops.cc
@@ -1818,7 +1818,11 @@ REGISTER_OP("_MklMaxPool")
.Input("input: T")
.Input("mkl_input: uint8")
.Output("output: T")
+#ifndef INTEL_MKL_DNN
.Output("workspace: T")
+#else
+ .Output("workspace: uint8")
+#endif
.Output("mkl_output: uint8")
.Output("mkl_workspace: uint8")
.SetShapeFn(shape_inference::MaxPoolShape)
@@ -1840,7 +1844,11 @@ REGISTER_OP("_MklMaxPoolGrad")
.Input("orig_input: T")
.Input("orig_output: T")
.Input("grad: T")
+#ifndef INTEL_MKL_DNN
.Input("workspace: T")
+#else
+ .Input("workspace: uint8")
+#endif
.Input("mkl_orig_input: uint8")
.Input("mkl_orig_output: uint8")
.Input("mkl_grad: uint8")