Error: GeneralizedDiceScore doesnt support batch size one

[asbjrnmunk]

When not limiting the amount of validation steps, one might encounter batches of size one if you have an uneven amount of validation cases. In this case GeneralizedDiceScore fails with the following error (when multilabel=False):

Input and target sizes: torch.Size([1, 1, 128, 128, 128]), torch.Size([1, 1, 128, 128, 128])

Error:

  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/metric.py", line 312, in forward
    self._forward_cache = self._forward_reduce_state_update(*args, **kwargs)
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/metric.py", line 381, in _forward_reduce_state_update
    self.update(*args, **kwargs)
  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/metric.py", line 493, in wrapped_func
    raise err
  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/metric.py", line 483, in wrapped_func
    update(*args, **kwargs)
  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/segmentation/generalized_dice.py", line 133, in update
    numerator, denominator = _generalized_dice_update(
                             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/functional/segmentation/generalized_dice.py", line 50, in _generalized_dice_update
    _check_same_shape(preds, target)
  File "/dtu/p1/asbmun/.cache/pypoetry/virtualenvs/virtualenvs/agave-i4r26evO-py3.11/lib/python3.11/site-packages/torchmetrics/utilities/checks.py", line 40, in _check_same_shape
    raise RuntimeError(
RuntimeError: Predictions and targets are expected to have the same shape, but got torch.Size([1, 2, 128, 128, 128]) and torch.Size([128, 2, 128, 128]).

This error does not appear for batch sizes >1. (2 and 4 tested).

[asbjrnmunk]

This workaround works for now, but should be fixed at some point:

     def validation_step(self, batch, batch_idx):
         inputs, target, file_path = batch["image"], batch["label"], batch["file_path"]

+        # Fix: Generalized Dice Score doesnt support batch size one
+        # so we skip validation if this is the case.
+        if inputs.size(0) == 1:
+            return
+
         output = self(inputs)