Project import generated by Copybara.

PiperOrigin-RevId: 219386646

Author: tf-transform-team

tensorflow_transform/analyzers.py

@@ -974,7 +974,8 @@ def initialize_local_state(self, tf_config=None):
     stamp_token = 0
 
     # Create a new session with a new graph for quantile ops.
-    self._session = tf.Session(graph=tf.Graph(), config=tf_config)
+    graph = tf.Graph()
+    self._session = tf.Session(graph=graph, config=tf_config)
     with self._session.graph.as_default():
       with self._session.as_default():
         self._qaccumulator = quantile_ops.QuantileAccumulator(
@@ -1025,6 +1026,8 @@ def initialize_local_state(self, tf_config=None):
         _, self._buckets_op = self._qaccumulator.get_buckets(
             stamp_token=stamp_token)
 
+        graph.finalize()
+
     # We generate an empty summary by calling self._flush_summary_op.
     # We cache this as some implementations may call create_accumulator for
     # every input, and it can be cached since it will always be the same and

tensorflow_transform/beam/impl.py

@@ -287,6 +287,7 @@ def __init__(self, saved_model_dir, input_schema, exclude_outputs,
                   saved_model_dir, {}))
         self.session.run(tf.global_variables_initializer())
         self.session.run(tf.tables_initializer())
+        graph.finalize()
 
         input_schema_keys = input_schema.column_schemas.keys()
         extra_input_keys = set(input_schema_keys).difference(inputs.keys())