So I tried what was suggested, and I’m still quite confused, since now I am getting the following error:
logit shape torch.Size([32, 224, 224])
label shape torch.Size([32, 224, 224])
Traceback (most recent call last):
  File "train_segnet.py", line 185, 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 916, in forward
    ignore_index=self.ignore_index, reduction=self.reduction)
  File "venv/lib/python3.6/site-packages/torch/nn/functional.py", line 2009, in cross_entropy
    return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
  File "venv/lib/python3.6/site-packages/torch/nn/functional.py", line 1848, in nll_loss
    out_size, target.size()))
ValueError: Expected target size (32, 224), got torch.Size([32, 224, 224])

However:

the network seems to output values in a different range (I am using DeepLabV3('resnet18', classes=1) which is from this repository here: https://github.com/qubvel/segmentation_models.pytorch 1

that was the reason I got everything placed in different channels (the repo uses that approach)
I’ve placed my data as you suggested, however that doesn’t seem to work for me right now.

One of the errors I am currently getting is:
/pytorch/aten/src/THCUNN/SpatialClassNLLCriterion.cu:104: void cunn_SpatialClassNLLCriterion_updateOutput_kernel(T *, T *, T *, long *, T *, int, int, int, int, int, long) [with T = float, AccumT = float]: block: [7,0,0], thread: [991,0,0] Assertiont >= 0 && t < n_classesfailed
Which makes sense because when I look at the network output/logits they are:
[ 0.3511,  0.3759,  0.4006,  ...,  0.2662,  0.2667,  0.2672],
          [ 0.3513,  0.3746,  0.3980,  ...,  0.3169,  0.3135,  0.3100],
          [ 0.3514,  0.3734,  0.3954,  ...,  0.3677,  0.3602,  0.3528],
          ...,
          [ 0.2786,  0.2707,  0.2628,  ...,  0.0179, -0.0218, -0.0616],
          [ 0.2405,  0.2318,  0.2231,  ..., -0.0140, -0.0691, -0.1243],
          [ 0.2024,  0.1928,  0.1833,  ..., -0.0458, -0.1164, -0.1870]]],


        [[[-0.1596, -0.0971, -0.0346,  ...,  0.2567,  0.2747,  0.2928],
          [-0.0786, -0.0282,  0.0221,  ...,  0.2614,  0.2682,  0.2750],
          [ 0.0025,  0.0406,  0.0787,  ...,  0.2662,  0.2617,  0.2572],
          ...,
          [ 0.2729,  0.2585,  0.2440,  ..., -0.1774, -0.2158, -0.2543],
          [ 0.2877,  0.2605,  0.2333,  ..., -0.2256, -0.2790, -0.3325],
          [ 0.3024,  0.2625,  0.2225,  ..., -0.2739, -0.3423, -0.4107]]],


        [[[ 0.0162,  0.0238,  0.0313,  ...,  0.0944,  0.0774,  0.0605],
          [ 0.0553,  0.0645,  0.0736,  ...,  0.1529,  0.1334,  0.1138],
          [ 0.0943,  0.1052,  0.1160,  ...,  0.2115,  0.1893,  0.1671],
          ...,
          [-0.0165,  0.0217,  0.0599,  ..., -0.0660, -0.0712, -0.0764],
          [-0.0538, -0.0219,  0.0100,  ..., -0.0507, -0.0527, -0.0547],
          [-0.0912, -0.0655, -0.0399,  ..., -0.0355, -0.0342, -0.0329]]],


        ...,

So, is there some transformation I should apply to the logits?
Or should I go back to dealing with 7-channel output, and then translate those, to a 1-dim tensor which can be used to calculate losses, in the shape you have suggested?
I can obviously go from 7-channel per class, to 1-channel with class index values, and then find a way to translate those so that the losses can be calculated, assuming that will work as per your example.