Suppose I have a model that takes as input a multi-level mask like a segmentation map where each pixel can take one of N > 2 classes.
My question is should the input shape be (batch_size, num_classes, height, width) or can I instead use only one channel to encode the class index with an input shape of (batch_size, 1, height, width).
With output shape its pretty clear that it is (batch_size, num_classes, height, width). But what about the input shape?
Please advice.