From 65d41488cb6e1057d122ab024bf48d4f86a6e48c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ricardo=20Monta=C3=B1ana?= <rmontanana@gmail.com>
Date: Thu, 29 Dec 2022 00:45:10 +0100
Subject: [PATCH] Fix AODE state_names

---
 bayesclass/clfs.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/bayesclass/clfs.py b/bayesclass/clfs.py
index 25d1c9d..c47adca 100644
--- a/bayesclass/clfs.py
+++ b/bayesclass/clfs.py
@@ -335,7 +335,7 @@ class AODE(BayesBase, BaseEnsemble):
 
         self.dag_ = None
 
-    def _train(self):
+    def _train(self, kwargs):
         """Build SPODE estimators (Super Parent One Dependent Estimator)"""
         self.models_ = []
         class_edges = [(self.class_name_, f) for f in self.features_]
@@ -353,6 +353,7 @@ class AODE(BayesBase, BaseEnsemble):
                 self.dataset_,
                 estimator=BayesianEstimator,
                 prior_type="K2",
+                state_names=kwargs["state_names"],
             )
             self.models_.append(model)