train_data = train_data.cache().prefetch(buffer_size=tf.data.AUTOTUNE) val_data = val_data.cache().prefetch(buffer_size=tf.data.AUTOTUNE)