def METHOD_NAME(
    self,
    sample: dict[str, Any],
    show_titles: bool = True,
    suptitle: Optional[str] = None,
) -> plt.Figure:
    """Plot a sample from the dataset.
    Args:
        sample: a sample returned by :meth:`RasterDataset.__getitem__`
        show_titles: flag indicating whether to show titles above each panel
        suptitle: optional string to use as a suptitle
    Returns:
        a matplotlib Figure with the rendered sample
    """
    mask = sample["mask"].squeeze()
    ncols = 1
    showing_predictions = "prediction" in sample
    if showing_predictions:
        prediction = sample["prediction"].squeeze()
        ncols = 2
    fig, axs = plt.subplots(nrows=1, ncols=ncols, figsize=(4 * ncols, 4))
    if showing_predictions:
        axs[0].imshow(mask)
        axs[0].axis("off")
        axs[1].imshow(prediction)
        axs[1].axis("off")
        if show_titles:
            axs[0].set_title("Mask")
            axs[1].set_title("Prediction")
    else:
        axs.imshow(mask)
        axs.axis("off")
        if show_titles:
            axs.set_title("Mask")
    if suptitle is not None:
        plt.suptitle(suptitle)
    return fig