def _train_and_eval_dataset_v1(problem_name, data_dir, train_shuffle_files,
                               eval_shuffle_files):
  """Return train and evaluation datasets, feature info and supervised keys."""
  with tf.device('cpu:0'):
    problem = t2t_problems().problem(problem_name)
    hparams = None
    if problem_name == 'video_bair_robot_pushing':
      hparams = problem.get_hparams()
      bair_robot_pushing_hparams(hparams)
    train_dataset = problem.dataset(
        tf_estimator.ModeKeys.TRAIN,
        data_dir,
        shuffle_files=train_shuffle_files,
        hparams=hparams)
    train_dataset = train_dataset.map(_select_features)
    eval_dataset = problem.dataset(
        tf_estimator.ModeKeys.EVAL,
        data_dir,
        shuffle_files=eval_shuffle_files,
        hparams=hparams)
    eval_dataset = eval_dataset.map(_select_features)
    # TODO(lukaszkaiser): remove this need for one example, just input_key.
    examples = list(tfds.as_numpy(train_dataset.take(1)))
  # We use 'inputs' as input except for purely auto-regressive tasks like
  # language models where 'targets' are used as input_key.
  input_key = 'inputs' if 'inputs' in examples[0] else 'targets'
  supervised_keys = ([input_key], ['targets'])
  return train_dataset, eval_dataset, supervised_keys
Return train and evaluation datasets, feature info and supervised keys.
trax/data/tf_inputs.py