def add_eos_to_output_features(dataset,
                               training,
                               output_features='targets',
                               eos=1):
  """Adds `EOS` to all features in `output_features`."""
  del training
  if not isinstance(output_features, (list, tuple)):
    output_features = [output_features]

  def add_eos(x):
    for output_feature in output_features:
      x[output_feature] = tf.concat([x[output_feature], [eos]], axis=0)
    return x

  return dataset.map(add_eos)
Adds `EOS` to all features in `output_features`.
trax/data/tf_inputs.py