Error for index tensor verification in GeneralizedDice if target is zeros

[juliendenize]

🐛 Bug

To compute segmentation mask GeneralizedDice metrics one can use index tensors as stated by the documentation.

However, if the mask is filled by zeros the conversion to one_hot label will not work in the following lines:

if (target.bool() != target).any():  # target is an index tensor
    target = torch.nn.functional.one_hot(target, num_classes=num_classes).movedim(-1, 1)

To Reproduce

Run the following code snippet:

import torch
from torchmetrics.segmentation import GeneralizedDiceScore

metric = GeneralizedDiceScore(num_classes=3)

target = torch.zeros((2, 3, 3), dtype=torch.int64)
predictions = torch.randint(0, 3, (2, 3, 3), dtype=torch.int64)

metric(predictions, target)

>>> RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1

Expected behavior

The expected behavior is to have the metric computed after the conversion of the target to a one-hot label.

Environment

• TorchMetrics version 1.4.0

Additional context

[github-actions]

Hi! thanks for your contribution!, great first issue!