def input_fn():
    """Input function for training and eval."""
    dataset = tf.data.experimental.make_batched_features_dataset(
        file_pattern=transformed_examples,
        batch_size=batch_size,
        features=tf_transform_output.transformed_feature_spec(),
        reader=tf.data.TFRecordDataset,
        shuffle=True)

    transformed_features = _make_inputs_dense(
        tf.compat.v1.data.make_one_shot_iterator(dataset).get_next()
    )

    # Extract features and label from the transformed tensors.
    # TODO(b/30367437): make transformed_labels a dict.
    transformed_labels = tf.where(
        tf.equal(transformed_features.pop(common.LABEL_KEY), 1))

    return transformed_features, transformed_labels[:, 1]
Input function for training and eval.
examples/census_example.py