import tensorflow as tf

            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():
            if tensor.dtype == tf.int64:
                example[name] = tf.cast(tensor, tf.int32)
            else:
                example[name] = tensor

