AUTOTUNE = tf.data.experimental.AUTOTUNE

train_dataset = train_dataset.cache()
train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)

        __
        __