Add the plot of probabilities after prediction