hausdorff_distance() throws IndexError for certain input shapes

import torchmetrics.functional.segmentation
import torch
import sys
print(torch.__version__)
print(torchmetrics.__version__)
print(sys.version)

def main(i):
    if i == 0:
        shape = (1, 1, 10, 10) # works
    elif i == 1:
        shape = (1, 1, 10, 11) # works
    else:
        shape = (1, 1, 11, 10) # throws error
    a = torch.randint(0, 2, size=shape)
    b = torch.randint(0, 2, size=shape)
    return torchmetrics.functional.segmentation.hausdorff_distance(a, b, num_classes=1)

print(f'{main(0)=}')
print(f'{main(1)=}')
print(f'{main(2)=}')

  File "/home/user/script.py", line 21, in <module>
    print(f'{main(2)=}')
             ~~~~^^^
  File "/home/user/script.py", line 17, in main
    return torchmetrics.functional.segmentation.hausdorff_distance(a, b, num_classes=1)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/hausdorff_distance.py", line 99, in hausdorff_distance
    dist = edge_surface_distance(
        preds=preds[b, c],
    ...<3 lines>...
        symmetric=not directed,
    )
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/utils.py", line 499, in edge_surface_distance
    surface_distance(edges_preds, edges_target, distance_metric=distance_metric, spacing=spacing),
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/utils.py", line 469, in surface_distance
    dis = distance_transform(~target, sampling=spacing, metric=distance_metric)
  File "/home/user/miniforge3/envs/general/lib/python3.13/site-packages/torchmetrics/functional/segmentation/utils.py", line 351, in distance_transform
    z[i1 * h + j1] = mindis
    ~^^^^^^^^^^^^^
IndexError: index 156 is out of bounds for dimension 0 with size 156

shell returned 1

• TorchMetrics version: 1.8.2
• Python & PyTorch Version: 3.13.3 (conda-forge), 2.6.0
• Any other relevant information such as OS (e.g., Linux): Ubuntu