diff --git a/src/mlpack/methods/naive_bayes/naive_bayes_classifier_impl.hpp b/src/mlpack/methods/naive_bayes/naive_bayes_classifier_impl.hpp index 9dbd7c533c..21829e3938 100644 --- a/src/mlpack/methods/naive_bayes/naive_bayes_classifier_impl.hpp +++ b/src/mlpack/methods/naive_bayes/naive_bayes_classifier_impl.hpp @@ -146,9 +146,11 @@ void NaiveBayesClassifier::Train(const VecType& point, arma::vec delta = point - means.col(label); means.col(label) += delta / probabilities[label]; - variances.col(label) *= (probabilities[label] - 1) / probabilities[label]; - variances.col(label) += (1 / probabilities[label]) * - (delta % (point - means.col(label))); + if (probabilities[label] > 2) + variances.col(label) *= (probabilities[label] - 2); + variances.col(label) += (delta % (point - means.col(label))); + if (probabilities[label] > 1) + variances.col(label) /= probabilities[label] - 1; trainingPoints++; probabilities /= trainingPoints;