Repository: Lightning-AI/torchmetrics

Issue:
CalibrationError bin_boundaries added to state_dict
Hi there,

The `bin_boundaries` of the `CalibrationError` metrics are added to the state_dict when model are saved (e.g., in PL: `trainer.save_checkpoint`). This creates an error when the state_dict is reloaded (e.g., in PL: `MyModel.load_from_checkpoint`) because the keys relative to the `bin_boundaries` are not available in the model.

cc @SkafteNicki 

https://github.com/PyTorchLightning/metrics/blob/d27f17107b7ee9c118f1dc88eaf2b58147ff4400/torchmetrics/classification/calibration_error.py#L84


Discussion:

Resolving patch:
diff --git a/torchmetrics/classification/binned_precision_recall.py b/torchmetrics/classification/binned_precision_recall.py
--- a/torchmetrics/classification/binned_precision_recall.py
+++ b/torchmetrics/classification/binned_precision_recall.py
@@ -124,14 +124,13 @@ def __init__(
         self.num_classes = num_classes
         if isinstance(thresholds, int):
             self.num_thresholds = thresholds
-            thresholds = torch.linspace(0, 1.0, thresholds)
-            self.register_buffer("thresholds", thresholds)
+            self.thresholds = torch.linspace(0, 1.0, thresholds)
+
         elif thresholds is not None:
             if not isinstance(thresholds, (list, Tensor)):
                 raise ValueError("Expected argument `thresholds` to either be an integer, list of floats or a tensor")
-            thresholds = torch.tensor(thresholds) if isinstance(thresholds, list) else thresholds
-            self.num_thresholds = thresholds.numel()
-            self.register_buffer("thresholds", thresholds)
+            self.thresholds = torch.tensor(thresholds) if isinstance(thresholds, list) else thresholds
+            self.num_thresholds = self.thresholds.numel()
 
         for name in ("TPs", "FPs", "FNs"):
             self.add_state(
diff --git a/torchmetrics/classification/calibration_error.py b/torchmetrics/classification/calibration_error.py
--- a/torchmetrics/classification/calibration_error.py
+++ b/torchmetrics/classification/calibration_error.py
@@ -75,7 +75,7 @@ def __init__(
         if not isinstance(n_bins, int) or n_bins <= 0:
             raise ValueError(f"Expected argument `n_bins` to be a int larger than 0 but got {n_bins}")
         self.n_bins = n_bins
-        self.register_buffer("bin_boundaries", torch.linspace(0, 1, n_bins + 1))
+        self.bin_boundaries = torch.linspace(0, 1, n_bins + 1)
         self.norm = norm
 
         self.add_state("confidences", [], dist_reduce_fx="cat")
@@ -102,4 +102,4 @@ def compute(self) -> Tensor:
         """
         confidences = dim_zero_cat(self.confidences)
         accuracies = dim_zero_cat(self.accuracies)
-        return _ce_compute(confidences, accuracies, self.bin_boundaries, norm=self.norm)
+        return _ce_compute(confidences, accuracies, self.bin_boundaries.to(self.device), norm=self.norm)
diff --git a/torchmetrics/wrappers/minmax.py b/torchmetrics/wrappers/minmax.py
--- a/torchmetrics/wrappers/minmax.py
+++ b/torchmetrics/wrappers/minmax.py
@@ -65,8 +65,8 @@ def __init__(
                 f"Expected base metric to be an instance of `torchmetrics.Metric` but received {base_metric}"
             )
         self._base_metric = base_metric
-        self.register_buffer("min_val", torch.tensor(float("inf")))
-        self.register_buffer("max_val", torch.tensor(float("-inf")))
+        self.min_val = torch.tensor(float("inf"))
+        self.max_val = torch.tensor(float("-inf"))
 
     def update(self, *args: Any, **kwargs: Any) -> None:  # type: ignore
         """Updates the underlying metric."""
@@ -83,8 +83,8 @@ def compute(self) -> Dict[str, Tensor]:  # type: ignore
             raise RuntimeError(
                 f"Returned value from base metric should be a scalar (int, float or tensor of size 1, but got {val}"
             )
-        self.max_val = val if self.max_val < val else self.max_val
-        self.min_val = val if self.min_val > val else self.min_val
+        self.max_val = val if self.max_val.to(val.device) < val else self.max_val.to(val.device)
+        self.min_val = val if self.min_val.to(val.device) > val else self.min_val.to(val.device)
         return {"raw": val, "max": self.max_val, "min": self.min_val}
 
     def reset(self) -> None:


Test patch:
diff --git a/tests/helpers/testers.py b/tests/helpers/testers.py
--- a/tests/helpers/testers.py
+++ b/tests/helpers/testers.py
@@ -121,6 +121,7 @@ def _class_test(
     device: str = "cpu",
     fragment_kwargs: bool = False,
     check_scriptable: bool = True,
+    check_state_dict: bool = True,
     **kwargs_update: Any,
 ):
     """Utility function doing the actual comparison between class metric and reference metric.
@@ -213,6 +214,10 @@ def _class_test(
     # check that metrics are hashable
     assert hash(metric)
 
+    # assert that state dict is empty
+    if check_state_dict:
+        assert metric.state_dict() == {}
+
     # check on all batches on all ranks
     result = metric.compute()
     if isinstance(result, dict):
diff --git a/tests/image/test_lpips.py b/tests/image/test_lpips.py
--- a/tests/image/test_lpips.py
+++ b/tests/image/test_lpips.py
@@ -59,6 +59,7 @@ def test_lpips(self, net_type, ddp):
             sk_metric=partial(_compare_fn, net_type=net_type),
             dist_sync_on_step=False,
             check_scriptable=False,
+            check_state_dict=False,
             metric_args={"net_type": net_type},
         )