def plot_to_image(figure):
    """Converts the matplotlib plot specified by 'figure' to a PNG image and
    returns it. The supplied figure is closed and inaccessible after this call."""
    # Save the plot to a image in memory.
    buf = io.BytesIO()
    plt.savefig(buf, format="png")
    # Closing the figure prevents it from being displayed directly inside
    # the notebook.
    plt.close(figure)
    buf.seek(0)
    # Convert PNG buffer to PyTorch tensor
    image = Image.open(buf)
    # Convert PIL Image to a NumPy array
    numpy_array = np.array(image)
    # Convert the NumPy array to a PyTorch tensor
    torch_image = torch.from_numpy(numpy_array)
    # # Ensure the data type and normalize if needed (optional)
    # torch_image = torch_image.float() / 255.0  # Normalize to [0, 1] if the image has pixel values in [0, 255]
    # Convert to CHW format by rearranging the dimensions
    torch_image_chw = torch_image.permute(2, 0, 1)

    return torch_image_chw
Converts the matplotlib plot specified by 'figure' to a PNG image and
    returns it. The supplied figure is closed and inaccessible after this call.
smoe/utils/visualization/visualize.py