aboutsummaryrefslogtreecommitdiffhomepage
diff options
context:
space:
mode:
authorGravatar Alexandre Passos <apassos@google.com>2018-07-27 09:31:48 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-07-27 09:35:35 -0700
commite6138e96def4b6432ad930cfb34b3806c5f9fb8b (patch)
treebe2d260d79c11a0500811bd05877824b7c00aac8
parent194e724cf9581ae7c3688b3942a66321d9fb514d (diff)
Properly call PrepareToUpdateVariable in resource strided slice assign.
PiperOrigin-RevId: 206327963
-rw-r--r--tensorflow/core/kernels/BUILD1
-rw-r--r--tensorflow/core/kernels/strided_slice_op.cc3
2 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 2cb54bd973..65a7f8ccf3 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -124,6 +124,7 @@ tf_kernel_library(
":bounds_check",
":dense_update_functor",
":ops_util",
+ ":training_op_helpers",
":variable_ops",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
diff --git a/tensorflow/core/kernels/strided_slice_op.cc b/tensorflow/core/kernels/strided_slice_op.cc
index 1e3e92a68a..c08ff47583 100644
--- a/tensorflow/core/kernels/strided_slice_op.cc
+++ b/tensorflow/core/kernels/strided_slice_op.cc
@@ -32,6 +32,7 @@ limitations under the License.
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/kernels/bounds_check.h"
#include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/kernels/training_op_helpers.h"
#include "tensorflow/core/kernels/variable_ops.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
@@ -304,6 +305,8 @@ class StridedSliceAssignOp : public OpKernel {
Var* v;
OP_REQUIRES_OK(context,
LookupResource(context, HandleFromInput(context, 0), &v));
+ OP_REQUIRES_OK(context,
+ PrepareToUpdateVariable<Device, T>(context, v->tensor()));
old_lhs = *v->tensor();
OP_REQUIRES(context, old_lhs.dtype() == DataTypeToEnum<T>::value,
errors::InvalidArgument(