import tensorflow as tf


    def _input_fn_builder(self, input_file, is_training):
        """Creates an `input_fn` closure to be passed to TPUEstimator."""

        def input_fn(params):
            """The actual input function."""
            d = tf.data.TFRecordDataset(input_file)

            if is_training:
                d = d.repeat()
                d = d.shuffle(buffer_size=100)
            return d.apply(
