fixes wasserstein_dice shape
@@ -313,6 +313,7 @@ def wasserstein_disagreement_map(
# pixel-wise Wassertein distance (W) between flat_pred_proba and flat_labels
# wrt the distance matrix on the label space M
n_classes = prediction.shape[1].value
+ ground_truth.set_shape(prediction.shape)
unstack_labels = tf.unstack(ground_truth, axis=-1)
unstack_labels = tf.cast(unstack_labels, dtype=tf.float64)
unstack_pred = tf.unstack(prediction, axis=-1)
