Thanks for your help, I fixed the issue by measuring dice score for each image in the validation batch and then averaging, instead of measuring dice once on the entire batch.
On the other hand, I found that training converged much faster and was much more stable when using my original ‘batch dice loss’ instead (i.e. same as what is posted here: https://github.com/pytorch/pytorch/issues/1249#issuecomment-305088398 11)