Hi,
I’m applying nn.NLLLoss () for segmentation problems.
However, I don’t understand how to extract the segmentation from the final layer. I’ll explain.

Given an encoderdecoder model, the final decoder level releases a dimension tensor (N, C, W, H)
to which softmax is applied, where N=minibach, C= number of class, WxH= image size. 
the tensor target has the following structure (N,W,H), wherein WxH class 0 contains the object to be segmented (a), while class 1 contains the object to be segmented (b), a + b = total segmentation.

now my doubt: from the output decoder of the decoder I want to extract my two objects to be segmented (my two classes) therefore strictly logic:
 output tesnor 1: (N,[0],W,H) – > object a
 output tensor 2: (N,[1],W,H) – > object b
to extract the elementation I will have to make a threshold or search for the value 0 or 1 in output tensor 1,2?
For example:
t1_seg = tensor_out_decoder[0]
t2_seg = tensor_out_decoder[1]
a = np.zeros(shape=(seg.shape[1],seg.shape[2]))
b = np.zeros(shape=(seg.shape[1],seg.shape[2]))
idx_a = t1_seg[:,:] > 0.5
idx_b = t2_seg[:,:] > 0.5
a[idx_myo] = 50
b[idx_rv] = 100
Best,
Nico