Uses a resource variable by default for the global step.

PiperOrigin-RevId: 200467580

1 files changed, 4 insertions, 3 deletions

diff --git a/tensorflow/python/saved_model/builder_impl.py b/tensorflow/python/saved_model/builder_impl.py
index e58be804c2..531da052ac 100644
--- a/tensorflow/python/saved_model/builder_impl.py
+++ b/tensorflow/python/saved_model/builder_impl.py
@@ -28,6 +28,7 @@ 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
@@ -178,11 +179,11 @@ 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`.
+      TypeError if Train op is not of type `Operation` or a Tensor.
     """
     if train_op is not None:
-      if (not isinstance(train_op, ops.Tensor) and
-          not isinstance(train_op, ops.Operation)):
+      if not (tensor_util.is_tensor(train_op) or
+              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)