Bug in the generalized Wasserstein Dice loss

[LucasFidon]

Describe the bug
The values returned by the loss are not correct for batch sizes larger than 1

To Reproduce
In the binary case and for M=[[1,0],[0,1]] the generalized Wasserstein Dice loss should be the same as the Dice loss.
For batch size > 1 it is not the case.

Expected behavior
See above. For this special case the GWDL and the DL should return the same values.

Additional context
I am the author of the generalized Wasserstein Dice loss paper and the bug has been already fixed in my git repository
https://github.com/LucasFidon/GeneralizedWassersteinDiceLoss
I will fix the bug in MONAI too and add unit tests.

[hongjie-qiu]

Hey! I noticed the same problem, found this existing open issue, and went ahead to put together a fix.
The root cause is that alpha_extended keeps shape (B, 1, S) after the torch.gather in _compute_generalized_true_positive and _compute_denominator, so when it's multiplied with wasserstein_distance_map (B, S) it silently broadcasts to (B, B, S) for any batch size > 1. The fix squeezes that extra dim after gather, matching the reference implementation. Also fixed the reduction="none" path which had a related reshape issue.

cc @KumoLiu would appreciate a review when you get a chance