Hi, I’m trying to build a segmentation model with 3 classes.
This is my way of action:
1.my output from the model is :(1,3,512,512)
2. softmax on channel dimenssion.
3. argmax on channel dimension.
4. getting (1,512,512) tensor, correct so far.
When I’m trying to enter to NLLLoss2d , I’m getting an error:
expected … (1,512) , and not (1,512,512)
So I moved to  NLLLoss , and tried to flatten my tensors into : (1,262144)
This also didn’t work, I’m getting:
"RuntimeError: multi-target not supported at c:\new-builder_2\win-wheel\pytorch\aten\src\thnn\generic/ClassNLLCriterion.c:21 "
What I’m missing? in Keras it’s so simple.
Please help me understand how to train segmentation models.
Thanks!