Shapes are incompatible for multi class classifier

I have a multi-class classifier, that takes inputs from a generator:

    def generate_train_data(path, x_shape):
        genres = {"hip-hop":0, "r&b":1, "pop":2, "jazz":3}
        genre_labels = to_categorical(list(genres.values()), num_classes=len(genres))
        # some processing  to create variables x and genre...
        # (mock values)
        x = np.zeros(x_shape)
        x = x[None, :, :, :]
        genre = "hip-hop"
        yield (x, genre_labels[genres[genre]])

The classifier is defined below:

    input_shape = (96, 84, 5)
    i = Input(shape=input_shape, name='encoder_input')
    cx = Conv2D(filters=8, kernel_size=3, strides=2, padding='same', activation='relu')(i)
    cx      = BatchNormalization()(cx)
    cx      = Conv2D(filters=16, kernel_size=3, strides=2, padding='same', activation='relu')(cx)
    cx      = BatchNormalization()(cx)
    x       = Flatten()(cx)
    x       = Dense(20, activation='relu')(x)
    x       = BatchNormalization()(x)
    x       = Dense(4, activation='softmax')(x)
    classifier = Model(i, x, name='genre_classifier')
    classifier.summary()
    classifier.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

However, when I try to fit the classifier:

classifier.fit(generate_train_data(path, input_shape), epochs=30, validation_data=generate_test_data(path, input_shape), verbose=verbosity)

I get the following error:

ValueError: in user code:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:806 train_function  *
    return step_function(self, iterator)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:796 step_function  **
    outputs = model.distribute_strategy.run(run_step, args=(data,))
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:1211 run
    return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2585 call_for_each_replica
    return self._call_for_each_replica(fn, args, kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2945 _call_for_each_replica
    return fn(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:789 run_step  **
    outputs = model.train_step(data)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:749 train_step
    y, y_pred, sample_weight, regularization_losses=self.losses)
/usr/local/lib/python3.6/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.6/dist-packages/tensorflow/python/keras/losses.py:149 __call__
    losses = ag_call(y_true, y_pred)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/losses.py:253 call  **
    return ag_fn(y_true, y_pred, **self._fn_kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py:201 wrapper
    return target(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/losses.py:1535 categorical_crossentropy
    return K.categorical_crossentropy(y_true, y_pred, from_logits=from_logits)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py:201 wrapper
    return target(*args, **kwargs)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/backend.py:4687 categorical_crossentropy
    target.shape.assert_is_compatible_with(output.shape)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/tensor_shape.py:1134 assert_is_compatible_with
    raise ValueError("Shapes %s and %s are incompatible" % (self, other))

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

The class label value returned by the generators is an array of length 4 so why is keras suggesting it is of size 1?

NOTE: This code is being run on Colab, tensorflow version 2.3. A mock version that reprocuces this error can be found on this Colab link: https://colab.research.google.com/drive/1SQZFspj3UOwP2ApIiaI2lvB2Z59bdVOk?usp=sharing

EDIT: added mock values in generate_train_data so that code can be reproducible