🐛 Fix #99 error in calibration plots

ENSTA-U2IS-AI · Jun 6, 2024 · acbd582 · acbd582
1 parent 63fb874
commit acbd582
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/tests/metrics/classification/test_calibration.py b/tests/metrics/classification/test_calibration.py
@@ -15,6 +15,7 @@ def test_plot_binary(self) -> None:
             torch.as_tensor([0, 0, 1, 1, 1]),
         )
         fig, ax = metric.plot()
+        metric.plot(ax=ax)
         assert isinstance(fig, plt.Figure)
         assert isinstance(ax, plt.Axes)
         assert ax.get_xlabel() == "Top-class Confidence (%)"

diff --git a/torch_uncertainty/metrics/classification/calibration_error.py b/torch_uncertainty/metrics/classification/calibration_error.py
@@ -60,7 +60,8 @@ def _ce_plot(self, ax: _AX_TYPE | None = None) -> _PLOT_OUT_TYPE:
     ax.set_xlim(0, 100)
     ax.set_ylim(0, 100)
     ax.set_aspect("equal", "box")
-    fig.tight_layout()
+    if fig is not None:
+        fig.tight_layout()
     return fig, ax