I’m not seeing a softmax usage besides this commented out line:
           # preds = torch.log_softmax(model(x), 1)
           # preds_labels = torch.argmax(preds, 1)

which would not be needed before applying argmax, so you could remove it.
Besides that the code looks good.
I would probably not recreate the criterion inside the DataLoader loop and just reuse it instead.
Also, I don’t know how the dice_loss etc. is implemented and you could check the loss and accuracy calculations with a known example just as a smoke test.