import tensorflow as tf

  Preprocessing for downsampled_imagenet 32x32 and 64x64 generation from
  http://arxiv.org/abs/1601.06759 (page 8).
  """
  del training

  def flatten_image(features):
    img = features['image']
    flat = tf.cast(tf.reshape(img, [-1]), tf.int64)

    new_features = {'image': flat}
    return new_features

  return dataset.map(flatten_image)

