Automated rollback of commit 568727eed199dba04e37f500265b50f96fed455e

PiperOrigin-RevId: 205875586

1 files changed, 2 insertions, 3 deletions

diff --git a/tensorflow/python/saved_model/builder_impl.py b/tensorflow/python/saved_model/builder_impl.py
index b67d0f2362..e58be804c2 100644
--- a/tensorflow/python/saved_model/builder_impl.py
+++ b/tensorflow/python/saved_model/builder_impl.py
@@ -28,7 +28,6 @@ from tensorflow.core.protobuf import saved_model_pb2
 from tensorflow.core.protobuf import saver_pb2
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor_util
 from tensorflow.python.lib.io import file_io
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import tf_logging
@@ -179,10 +178,10 @@ class SavedModelBuilder(object):
         stored as a collection with key TRAIN_OP_KEY, but not executed.
 
     Raises:
-      TypeError if Train op is not of type `Operation` or a Tensor.
+      TypeError if Train op is not of type `Operation`.
     """
     if train_op is not None:
-      if (not tensor_util.is_tensor(train_op) and
+      if (not isinstance(train_op, ops.Tensor) and
           not isinstance(train_op, ops.Operation)):
         raise TypeError("train_op needs to be a Tensor or Op: %r" % train_op)
       ops.add_to_collection(constants.TRAIN_OP_KEY, train_op)