I am using a custom model based on SegNet to do segmentation for 3 different objects. Model is returning a tensor of size “[5, 3, 120, 160] (batch size, channel, height, width)” and I have a label where each channel contain information of a particular object, something like this:
[[      [1, 0, 0, 0],
        [1, 1, 0, 1],
        [0, 0, 0, 0],
        [1, 0, 0, 0]],

       [[0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 1, 1, 0],
        [0, 0, 0, 1]],

       [[0, 0, 1, 1],
        [0, 0, 0, 0],
        [1, 0, 0, 1],
        [0, 1, 1, 0]]]

Label:  torch.Size([5, 3, 120, 160])
I am using CrossEntropy loss function to calculate the loss but I am getting this error “1only batches of spatial targets supported (non-empty 3D tensors) but got targets of size: : [5, 3, 120, 160]”.
As far as I can understand, I need to have a one-hot encoded vector for every pixel to calculate the loss. In my case, the label does represent one hot encoding (for a pixel, only one channel in the label will have the value ‘1’). May I know what can be done to resolve the issue?