import tensorflow as tf

  # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
  # So cast all int64 to int32.
  for name in list(example.keys()):
    t = example[name]
    if t.dtype == tf.int64:
      t = tf.to_int32(t)
    example[name] = t

  return example


def main(_):
  tf.logging.set_verbosity(tf.logging.INFO)

  if FLAGS.use_hvd:
    hvd.init()

    if FLAGS.reduce_log and (hvd.rank() != 0):
      tf.logging.set_verbosity(tf.logging.ERROR)

    FLAGS.output_dir = FLAGS.output_dir if hvd.rank() == 0 else os.path.join(FLAGS.output_dir, str(hvd.rank()))

  if not FLAGS.do_train and not FLAGS.do_eval and not FLAGS.do_train_eval:
    raise ValueError("At least one of `do_train` or `do_eval` must be True.")
