Hi @ptrblck, thanks for your guidance. I got the mask one-hot encoded i.e. binary image with each label across channels. So, currently my mask shape looks like [height, width, n_labels]. Now while training the model, this should work fine, but how do I get the output with the color coding ? Since input is binary i.e. 1’s and 0’s, the output will have n_labels channels but still the output will be binary only i.e. 1’s and 0’s.
So, how to get the color coded output ? Do I need to use an bitwise_and or bitwise_or with channels and colors on the output to get the mask as shown above in output ?
Thanks 
Your target should have the class dimension in dim1, such that its shape is [batch_size, nb_classes, height, width] for a multi-label segmentation.
To create color codes, you could map the predicted classes back to color codes. However, it depends on your use case how you would like to visualize overlapping classes. E.g. you could plot the “highest/lowest” class or even mix the colors.