train_dataset = (
    train_data
    .map(lambda img, label: image_augmentation(img, label), num_parallel_calls=tf.data.AUTOTUNE)  
    .cache()
    .prefetch(buffer_size=tf.data.AUTOTUNE) 
)

valid_dataset = (
    valid_data
    .map(lambda img, label: normalize_img(img, label), num_parallel_calls=tf.data.AUTOTUNE)  
    .cache()  
    .prefetch(buffer_size=tf.data.AUTOTUNE)  
)

# Optimization for test_data
test_dataset = (
    test_data
    .map(lambda img, label: normalize_img(img, label), num_parallel_calls=tf.data.AUTOTUNE) 
    .prefetch(buffer_size=tf.data.AUTOTUNE) 
)

        __
        __