I am training a 3D U-Net and am trying to implement a Dice loss with Tensorflow. I created the following function, which can return either the Dice score or the correspondin