ptrblck:

nn.CrossEntropyLoss expects an output in the shape [batch_size, nb_classes=2, depth, height, width] in your use case, and a target in the shape [batch_size, depth, height, width] containing values in the range [0, 1] .


I am sorry for there was a mistake with the output shape. My output shape is [4,2, 128,128,128]. The target shape is [4,1,128,128,128].
The outputs_soft is obtained as follows.
outputs_soft= F.softmax(outputs,dim=1)

And computing the dice_loss
 loss_seg_dice = dice_loss(outputs_soft[:,1,:,:,:],label_batch == 1)

does not give any error.