So I’ve reverted to using targets which are in shape of [batch, height, width] and their values are the class index (so no more channel per class).
Which AFAIK is what you have suggested. The loss function seems to work without hacking into any type of tensor or channels, simply with torch.argmax, however now my accuracy doesn’t work on a pixel-wise level, so I’m trying to implement it this way:
            outputs  = model(input)
            preds    = torch.argmax(outputs, dim=1)
            matching += (preds == labels).sum().item()
            correct  += (matching / 224 * 224)

            b_size  = labels.shape[0]
            total   += b_size

This seems to give a lower Top-1 score than before, but not by much (98% vs 99%).
So I am testing it again with the deployment platform, and after running the following:
        t_in   = self.tf(image)
        t_in   = torch.unsqueeze(t_in, 0)
        logits = self.net(t_in)
        torch.set_printoptions(profile="full")
        torch.set_printoptions(precision=5)
        #logits = logits.squeeze()
        print(logits.shape)
        print(logits)
        preds  = torch.argmax(logits, dim=1)
        print(preds.shape)
        print(preds)
        return self.extract(image, logits)

I noticed that the pred now contains the pixel class index correctly for all the classes present in the document.
So, THANK YOU, I owe you a few rounds of beer, your help is greatly appreciated!!!