train_dataset = (train_dataset.map(lambda img,label : image_augmentation(img,label))
                  .prefetch(tf.data.AUTOTUNE))
valid_dataset = (valid_dataset.map(lambda img,label:normalize_img(img,label))
              .prefetch(tf.data.AUTOTUNE))
test_data = (
            test_dataset.map(lambda img,label:normalize_img(img,label))

            .prefetch(tf.data.AUTOTUNE)
)

        __
        __