From 0b897ccbe38b2838561e1f455742da978e24396c Mon Sep 17 00:00:00 2001 From: starryyu Date: Wed, 24 Jun 2026 09:26:43 +0800 Subject: [PATCH] fix: clarify threshold error message --- .../metrics/functional.py | 2 +- tests/test_metrics.py | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 1 deletion(-) create mode 100644 tests/test_metrics.py diff --git a/segmentation_models_pytorch/metrics/functional.py b/segmentation_models_pytorch/metrics/functional.py index 5fd75cad5..d20a77a8b 100644 --- a/segmentation_models_pytorch/metrics/functional.py +++ b/segmentation_models_pytorch/metrics/functional.py @@ -120,7 +120,7 @@ def get_stats( if torch.is_floating_point(output) and threshold is None: raise ValueError( - f"Output should be one of the integer types if ``threshold`` is not None, got {output.dtype}." + f"Output should be one of the integer types if ``threshold`` is None, got {output.dtype}." ) if torch.is_floating_point(output) and mode == "multiclass": diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 000000000..6cdb99bf6 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,17 @@ +import pytest +import torch + +from segmentation_models_pytorch.metrics import get_stats + + +def test_get_stats_explains_float_output_requires_threshold() -> None: + output = torch.tensor([[[0.1, 0.8], [0.3, 0.7]]], dtype=torch.float32) + target = torch.tensor([[[0, 1], [0, 1]]], dtype=torch.long) + + with pytest.raises(ValueError) as error: + get_stats(output, target, mode="binary", threshold=None) + + assert str(error.value) == ( + "Output should be one of the integer types if ``threshold`` is None, " + "got torch.float32." + )