Loss function for segmentation models

As I mentioned, it didn’t work for me, I’m getting :
“ValueError: Expected target size (1, 512), got torch.Size([1, 512, 512])”

Another attempt I tried, I did exactly as I did in keras and did “to_categorial” like operator, now my target (mask) is (1,3,512,512).

so this is what I did:

outputs = model(inputs)
loss = criterion(F.log_softmax(outputs,1), masks.type(torch.LongTensor))

but I got this error:
RuntimeError: invalid argument 3: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4 at c:\new-builder_2\win-wheel\pytorch\aten\src\thnn\generic/SpatialClassNLLCriterion.c:59

Note: the size of the masks is :(1,3,512,512), and the outputs size is also:(1,3,512,512)