Hi, depending on the problem we’ve had good luck with CrossEntropyLoss as well as Dice.
Assuming you have your model producing a two-channel probability map, you should create a 1D view of the 3D images and the respective mask this way:
criterion = nn.CrossEntropyLoss()
output = model(input)
output = output.permute(0,2,3,4,1).contiguous()
output = output.view(output.numel() // 2, 2)
mask = mask.view(-1)
loss = criterion(output, mask)

For the Dice index, here’s a 2D implementation. I can’t send you the 3D version right now as I’m on the go: https://github.com/rogertrullo/pytorch/blob/rogertrullo-dice_loss/torch/nn/functional.py#L708 190