diff options
author | 2018-07-27 09:31:48 -0700 | |
---|---|---|
committer | 2018-07-27 09:35:35 -0700 | |
commit | e6138e96def4b6432ad930cfb34b3806c5f9fb8b (patch) | |
tree | be2d260d79c11a0500811bd05877824b7c00aac8 | |
parent | 194e724cf9581ae7c3688b3942a66321d9fb514d (diff) |
Properly call PrepareToUpdateVariable in resource strided slice assign.
PiperOrigin-RevId: 206327963
-rw-r--r-- | tensorflow/core/kernels/BUILD | 1 | ||||
-rw-r--r-- | tensorflow/core/kernels/strided_slice_op.cc | 3 |
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( |