Output a tensor for reference metric

Lightning-AI · Sep 24, 2024 · 8541d97 · 8541d97
1 parent c7bdce9
commit 8541d97
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/unittests/segmentation/test_hausdorff_distance.py b/tests/unittests/segmentation/test_hausdorff_distance.py
@@ -48,8 +48,9 @@
 
 
 # Wrapper that converts to numpy to avoid Torch-to-numpy functional issues
-def torch_skimage_hausdorff_distance(p: torch.Tensor, t: torch.Tensor) -> float:
-    return skimage_hausdorff_distance(p.numpy(), t.numpy())
+def torch_skimage_hausdorff_distance(p: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    out = skimage_hausdorff_distance(p.numpy(), t.numpy())
+    return torch.tensor([out])
 
 
 @pytest.mark.parametrize(