Dear:
The problem is that, when i implement the semantic segmetation, the output shape of the semantic segmentation network is [batch_size, num_classes, width, height],it is encoded like one_hot,      but i just want to directly get the output which shape is [batch_size, width, height],in the output each pixel has the value 0,1,2…(indicate the predicted class index), what can i do to get the output’s shape like [batch_size, width, height] (without postprocess like argmax() ).
please help me, thanks all.