Incompatible shapes (None, 1) and (None, 5) with Keras VGGFace Finetuning

Categories to learn and predict:

df.race.unique()
array(['0', '1', '3', '2', '4'], dtype=object)

Data:

train_generator = image_gen.flow_from_dataframe(
    df_train,
    x_col=img_name,
    y_col=race,
    directory=str(data_folder),
    class_mode=sparse,
    target_size=(IMAGE_SIZE, IMAGE_SIZE),
    batch_size=BATCH_SIZE,
    shuffle=True,
)

val_generator = image_gen.flow_from_dataframe(
    df_val,
    x_col=img_name,
    y_col=race,
    directory=str(data_folder),
    class_mode=sparse,
    target_size=(IMAGE_SIZE, IMAGE_SIZE),
    batch_size=BATCH_SIZE,
    shuffle=False,
)

Model load and fit:

vggface_model = load_model(resnet50face.h5)
base_model = tf.keras.Model([vggface_model.input], vggface_model.get_layer(flatten_1).output)
base_model.trainable = False

last_layer = base_model.get_layer('avg_pool').output
hidden_layer = Flatten(name='flatten')(last_layer)
out_layer = Dense(5, activation='softmax', name='gender_classifier')(hidden_layer)
custom_base_model = tf.keras.Model(base_model.input, out_layer)

custom_base_model.compile(
              optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001),
              loss=categorical_crossentropy,
              metrics=['accuracy'])

history = custom_base_model.fit(
    x=train_generator, 
    validation_data=val_generator, 
    steps_per_epoch=20, 
    epochs=40)

Error message:

ValueError: in user code:

    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:855 train_function  *
        return step_function(self, iterator)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:845 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:1285 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:2833 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:3608 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:838 run_step  **
        outputs = model.train_step(data)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:797 train_step
        y, y_pred, sample_weight, regularization_losses=self.losses)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/compile_utils.py:204 __call__
        loss_value = loss_obj(y_t, y_p, sample_weight=sw)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/losses.py:155 __call__
        losses = call_fn(y_true, y_pred)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/losses.py:259 call  **
        return ag_fn(y_true, y_pred, **self._fn_kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:206 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/losses.py:1644 categorical_crossentropy
        y_true, y_pred, from_logits=from_logits)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py:206 wrapper
        return target(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/backend.py:4862 categorical_crossentropy
        target.shape.assert_is_compatible_with(output.shape)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/tensor_shape.py:1161 assert_is_compatible_with
        raise ValueError(Shapes %s and %s are incompatible % (self, other))

    ValueError: Shapes (None, 1) and (None, 5) are incompatible