Return train and eval batches with input name and shape.

def _train_and_eval_batches(dataset, data_dir, input_name, num_devices):
  """Return train and eval batches with input name and shape."""
  (train_data, eval_data, features_info, keys) = train_and_eval_dataset(
      dataset, data_dir)
  input_names, target_names = keys[0], keys[1]
  train_batches = shuffle_and_batch_data(
      train_data, target_names, features_info, training=True,
      num_devices=num_devices)
  train_eval_batches = shuffle_and_batch_data(  # Data for eval-on-train.
      train_data, target_names, features_info, training=False,
      num_devices=num_devices)
  eval_batches = shuffle_and_batch_data(
      eval_data, target_names, features_info, training=False,
      num_devices=num_devices)
  input_name = input_name or input_names[0]
  input_shape = features_info[input_name].shape
  return (train_batches, train_eval_batches, eval_batches,
          input_name, list(input_shape))