def wrapper(fn):
    """Decorator that runs decorated function in graph mode."""
    def _run_graph(*inputs):
      with context.graph_mode():  # pylint: disable=missing-docstring
        assert len(input_signature) == len(inputs)
        placeholders = list(map(_make_placeholder, input_signature))
        output_tensor = fn(*placeholders)
        with tf.compat.v1.Session() as sess:
          sess.run(tf.compat.v1.tables_initializer())
          return sess.run(output_tensor,
                          feed_dict=dict(zip(placeholders, inputs)))

    return _run_graph
Decorator that runs decorated function in graph mode.
tensorflow_transform/test_case.py