Hello,
I’ve read quite a few relevant topics here on discuss.pytorch.org such as:

Loss function for segmentation models 19
Convert pixel wise class tensor to image segmentation 7
FCN Implementation : Loss Function 5

I’ve tried with CrossEntropyLoss but it comes with problems I don’t know how to easily overcome.
So I’m now trying to use nn.NLLLoss with pytorch 1.3 after running the network logits through torch.nn.functional.log_softmax
The way the dataset is organised, is that each input has a 7-channel tensor, where each channel is a class, and then the dims (224x224) are the pixels for that class.
So because I can have multiple classes in an image, I want them all indexed, across their respective channels.
I’ve followed the instructions from the above topics, but I keep getting the same error below:
logits: torch.Size([32, 7, 224, 224])
target: torch.Size([32, 7, 224, 224])
Traceback (most recent call last):
  File "train_segnet.py", line 183, in <module>
    loss   = criterion(logits, labels.cuda().long())
  File "venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 541, in __call__
    result = self.forward(*input, **kwargs)
  File "venv/lib/python3.6/site-packages/torch/nn/modules/loss.py", line 204, in forward
    return F.nll_loss(input, target, weight=self.weight, ignore_index=self.ignore_index, reduction=self.reduction)
  File "venv/lib/python3.6/site-packages/torch/nn/functional.py", line 1840, in nll_loss
    ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: 1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size: : [32, 7, 224, 224]

I thought that NLLoss is capable of dealing with multiple dimension tensors?
Should I squeeze them all across a single axis? How do I solve this issue so that the network
can learn to produce those tensors as output?