Hi, I’m trying to use MultiMarginLoss and got the error ‘inconsistent target size’

I’m working on a segmentation task. I have 6 labels, and I want to classify each pixel.

The output of my network is NxCxWxH with batch size N = 1, number classes C=6, width and height of the image = 512x512 pixels. The size of the label tensor is 1x512x512.

In the documentation it’s written that the input to MultiMarginLoss has to be a 2D tensor and the target a 1D tensor (I had to debug also for the data type…I don’t understand why the input has to be a FloatTensor, while the target a LongTensor… anyways). I tried the following: (I give also class weights as input to the loss criterion)

- take the argmax of the output, getting a tensor of 1x1x512x512, using view (I called .contiguous() before) I get a tensor of size 1x262144 for the input to MultiMarginLoss. Using view the target becomes a 1D tensor of size 262144.
- since my batch size is 1, I squeezed that without using argmax, I get a 2D tensor of 6x262144 with each row containing the probability of each class and the columns span every pixel. The target size is 1D 262144.
- I read about MultiLabelMarginLoss where it says that the input can be both 1D or both 2D… So I tried to make the target the same size as the input, i.e. 6x262144

I tried all these and get the same error… does anyone have any suggestion?