Hi @ptrblck
I tried to use nn.CrossEntropyLoss() initially like so, with image set to float() and mask set to long() type:
    # loss function
    loss_fn = nn.CrossEntropyLoss() 

but it gave me this error:
  File "/tool/python/conda/env/gis36/lib/python3.6/site-packages/torch/nn/functional.py", line 1334, in nll_loss
    return torch._C._nn.nll_loss2d(input, target, weight, size_average, ignore_index, reduce)
RuntimeError: invalid argument 1: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4 at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THCUNN/generic/SpatialClassNLLCriterion.cu:14

I attempted to write my own cross entropy loss function, thinking that I had to, because of the shape of my label outputs.
Is there a way to adapt and use the existing nn.CrossEntropyLoss() to work with the shape of my prediction and labels [batch, channels, height, width] ?
I think I’m missing something quite fundamental here. Is it okay to have my label masks with 10-channels, with each channel corresponding to a class (e.g. buildings)?
In your code snippet, labels is of shape [batch, height, width], whereas my labels are of shape [batch, channel, height, width]:
labels = torch.empty(4, 24, 24, dtype=torch.long).random_(n_class)