Keras categoricalcrossentropy for multi channel outuput format (NHWC)

Question

I´m building a segmentation model, and my output has the form of a image, where the channels are the categories, it means, suppose I have 5 categories, and an image of 4x4,