Fix to in the plotting of Naive Bayes