import tensorflow as tf

    aug_images, aug_generated = augment(images), augment(generated)

    # concat all images
    all_images = tf.concat([images, generated, aug_images, aug_generated], 0)

    if self.conditional:
