Multi-class segmentation

[Khola-naseem]

Hi, i am working on a dataset in which I have three classes(background, object1 and object2). I configured the model as described but still, the model is producing a binary mask. please tell me if I have to set any other option for multi-class segmentation

[haoyudong-97]

Can you provide more details? If you set args.num_cls to the right value, the output should have size num_cls, H, W and we have the final prediction as "pred.argmax(dim=1)", which means the output should be able to predict multiple classes.