Hi,
If this is just the cross entropy loss for each pixel independently, then you can use the existing cross entropy provided by pytorch.
The pytorch function only accepts input of size (batch_dim, n_classes).
So if your output is of size (batch, height, width, n_classes), you can use .view(batch * height * width, n_classes) before giving it to the cross entropy function (considering each pixel as a different batch element) to achieve what you want.