After researching for a while and using different resources I found out that the issue is with the way that you import the pre-trained model and not with any of the explained steps in the error.

So my code went from:

vit_model = hub.KerasLayer("https://tfhub.dev/sayakpaul/vit_b16_fe/1", trainable=False)

To:

vit_model = tf.keras.applications.VGG16(include_top=False, input_shape=(224, 224, 3))