def train(dataset, estimator, num_train_steps):
  """Run one training loop over given TFRecords file."""
  if FLAGS.profile_model:
    hooks = [
        tf.train.ProfilerHook(
            output_dir=estimator.model_dir, save_secs=100, show_memory=False)
    ]
    tf.logging.info("Saving profiling output to %s", estimator.model_dir)
  else:
    hooks = None
  estimator.train(
      input_fn=dataset.input_fn, max_steps=num_train_steps, hooks=hooks)
Run one training loop over given TFRecords file.
language/labs/drkit/run_multihop_follow.py