From 46e1183f368c7061b874dd3103714fec68348433 Mon Sep 17 00:00:00 2001 From: Nicki Skafte Detlefsen Date: Tue, 23 Apr 2024 13:31:21 +0200 Subject: [PATCH] Update src/torchmetrics/functional/segmentation/utils.py Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com> --- src/torchmetrics/functional/segmentation/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/torchmetrics/functional/segmentation/utils.py b/src/torchmetrics/functional/segmentation/utils.py index e8427a69326..6c2fed92df2 100644 --- a/src/torchmetrics/functional/segmentation/utils.py +++ b/src/torchmetrics/functional/segmentation/utils.py @@ -25,7 +25,7 @@ def _ignore_background(preds: Tensor, target: Tensor) -> Tuple[Tensor, Tensor]: - """Ignore the background class in the computation.""" + """Ignore the background class in the computation assuming it is the first, index 0.""" preds = preds[:, 1:] if preds.shape[1] > 1 else preds target = target[:, 1:] if target.shape[1] > 1 else target return preds, target