def plot(
        self, suggest: bool = False, show: bool = False, ax: Optional["Axes"] = None
    ) -> Optional[Union["plt.Figure", "plt.SubFigure"]]:
        """Plot results from lr_find run
        Args:
            suggest: if True, will mark suggested lr to use with a red point

            show: if True, will show figure

            ax: Axes object to which the plot is to be drawn. If not provided, a new figure is created.
        """
        if not _MATPLOTLIB_AVAILABLE:
            raise MisconfigurationException(
                "To use the `plot` method, you must have Matplotlib installed."
                " Install it by running `pip install -U matplotlib`."
            )
        import matplotlib.pyplot as plt

        lrs = self.results["lr"]
        losses = self.results["loss"]

        fig: Optional[Union[plt.Figure, plt.SubFigure]]
        if ax is None:
            fig, ax = plt.subplots()
        else:
            fig = ax.figure

        # Plot loss as a function of the learning rate
        ax.plot(lrs, losses)
        if self.mode == "exponential":
            ax.set_xscale("log")
        ax.set_xlabel("Learning rate")
        ax.set_ylabel("Loss")

        if suggest:
            _ = self.suggestion()
            if self._optimal_idx:
                ax.plot(lrs[self._optimal_idx], losses[self._optimal_idx], markersize=10, marker="o", color="red")

        if show:
            plt.show()

        return fig
Plot results from lr_find run
        Args:
            suggest: if True, will mark suggested lr to use with a red point

            show: if True, will show figure

            ax: Axes object to which the plot is to be drawn. If not provided, a new figure is created.
        
src/lightning/pytorch/tuner/lr_finder.py