def run_in_graph_mode_only(func=None, config=None, use_gpu=True):
  """Runs a test in graph mode only, when eager is enabled by default."""
  def decorator(f):
    """Decorator for a method."""
    def decorated(self, *args, **kwargs):
      """Run the decorated test method."""
      self.tearDown()
      # Run in graph mode block
      with tf.Graph().as_default():
        self.setUp()
        with self.test_session(use_gpu=use_gpu, config=config):
          f(self, *args, **kwargs)

    return decorated

  if func is not None:
    return decorator(func)

  return decorator
Runs a test in graph mode only, when eager is enabled by default.
tensor2tensor/utils/test_utils.py