def process(self, batch, saved_model_dir):
    """Runs the given graph to realize the outputs.

    Runs the graph in a TF session for computing the output values of the
    `Tensor`s, `SparseTensor`s, or `RaggedTensor`s, given an input row of data
    (input `Tensor`s, `SparseTensor`s, or `RaggedTensor`s).

    Args:
      batch: the batch of elements being processed by the DoFn
      saved_model_dir: Directory containing saved model.

    Yields:
      A representation of output features as a dict mapping keys (logical column
      names) to values.
    """
    if not hasattr(self, '_graph_state'):
      # If available, acquire will return a cached _GraphStateCommon, since
      # calling _make_graph_state is expensive.
      self._graph_state = self._shared_graph_state_handle.acquire(
          lambda: self._make_graph_state(saved_model_dir))

    # This should remain true throughout the lifetime of this DoFn, regardless
    # of whether or not self._graph_state was cached.
    assert self._graph_state.saved_model_dir == saved_model_dir

    yield self._handle_batch(batch)
Runs the given graph to realize the outputs.

    Runs the graph in a TF session for computing the output values of the
    `Tensor`s, `SparseTensor`s, or `RaggedTensor`s, given an input row of data
    (input `Tensor`s, `SparseTensor`s, or `RaggedTensor`s).

    Args:
      batch: the batch of elements being processed by the DoFn
      saved_model_dir: Directory containing saved model.

    Yields:
      A representation of output features as a dict mapping keys (logical column
      names) to values.
    
tensorflow_transform/beam/impl.py