You don’t need to apply a threshold for these shapes and use case.
The dice_loss should probably get the predicted class labels via:
preds = torch.argmax(outputs, 1)

Could you remove the softmax and pass preds to dice_loss?