import tensorflow as tf

            features[spec.name] = feature
        return tf.train.Example(features=tf.train.Features(feature=features))

    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(
                tf.contrib.data.map_and_batch(
                    self._decode_tfrecord, batch_size=params["batch_size"], drop_remainder=True
                )
            )

        return input_fn

    def _decode_tfrecord(self, record):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, self._name_to_feature_config)
        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.

        for name, tensor in example.items():
