Fix AttributeError: module 'matplotlib' has no attribute 'axes' (#1955)

Co-authored-by: Nicki Skafte Detlefsen <[email protected]>

Author: kn

src/torchmetrics/utilities/plot.py

@@ -23,6 +23,7 @@
 
 if _MATPLOTLIB_AVAILABLE:
     import matplotlib
+    import matplotlib.axes
     import matplotlib.pyplot as plt
 
     _PLOT_OUT_TYPE = Tuple[plt.Figure, Union[matplotlib.axes.Axes, np.ndarray]]