I see now why this has been so confusing. I’ve run your example and realised what’s going on.
The actual logits are of the shape of [batch, class_channel, height, width] which is what I’ve had all along. The labels/targets/ideals are in the shape of [batch, height, width]. That is the confusing part, and I didn’t really get why it’s like this, until I saw what was happening at the logits:

if they didn’t have a separate channel for the class or label I suspect there is no way to transform their raw/logit value to a class representation? Which is the problem I’ve run into twice so far.
So I transformed the labels into a 1-dim tensor, where each pixel represents a class index:

So I wrote this:
def segnet_label_translate(labels, num_classes):
    labels = labels.cuda().long()
    cuda0  = torch.device('cuda:0')
    b_size = labels.shape[0]
    tensor = torch.zeros([b_size, 224, 224], dtype=torch.long, device=cuda0)
    zeros  = torch.zeros([224, 224], dtype=torch.long, device=cuda0)

    for batch in range(b_size):
        for i in range(num_classes):
            ideal    = labels[batch][i]
            label    = i  + 1
            values   = torch.full([224, 224], label, dtype=torch.long, device=cuda0)
            filtered = torch.where((ideal != 0), values, zeros)
            tensor[batch] = filtered

    return tensor

I suspect there are easier of faster ways to write the above^.
That seems to work when I remove the line which does label = i + 1 otherwise I get the assertion triggered. However, what is unclear is why would I use class zero.
Then I remembered reading somewhere that class 0 in segmentation is always the background or some other less useful class?
The above works if I don’t do label = i + 1 but for me class 0 is something valid that I look for, so I suppose I’ll have to account for that.
Also, should I be applying a softmax or any other transformation to the logits, or is there no need for it?