ptrblck:

Could you post the custom loss and explain what your model outputs?
Since you are using nn.CrossEntropyLoss , I assume your model returns an output in the shape [batch_size, 2, height, width] ? If that’s the case, you should not apply any threshold, but get the predicted classes via preds = torch.argmax(output, 1) .


My loss function is as follows.
             outputs = net(volume_batch)
            
            loss_seg= F.cross_entropy(outputs,label_batch.long())
            outputs_soft= F.softmax(outputs,dim=1)
            loss_seg_dice = dice_loss(outputs_soft[:,1,:,:,:],label_batch == 1)
            loss= loss_seg + loss_seg_dice




 ptrblck:

If your model returns logits of the shape [batch_size, 1, height, width]


The shape of my model output is [4,1,128,128,128]. It is a 3D model and I have set 4 as batch size.