diff --git a/src/BayesNet/Network.cc b/src/BayesNet/Network.cc
index 6248ab5..3f304a4 100644
--- a/src/BayesNet/Network.cc
+++ b/src/BayesNet/Network.cc
@@ -325,17 +325,17 @@ namespace bayesnet {
         vector<thread> threads;
         mutex mtx;
         for (int i = 0; i < classNumStates; ++i) {
-            threads.emplace_back([this, &result, &evidence, i, &mtx]() {
+//            threads.emplace_back([this, &result, &evidence, i, &mtx]() {
                 auto completeEvidence = map<string, int>(evidence);
                 completeEvidence[getClassName()] = i;
                 double factor = computeFactor(completeEvidence);
-                lock_guard<mutex> lock(mtx);
+//                lock_guard<mutex> lock(mtx);
                 result[i] = factor;
-                });
-        }
-        for (auto& thread : threads) {
-            thread.join();
+//                });
         }
+//        for (auto& thread : threads) {
+//            thread.join();
+//        }
         // Normalize result
         double sum = accumulate(result.begin(), result.end(), 0.0);
         transform(result.begin(), result.end(), result.begin(), [sum](double& value) { return value / sum; });
diff --git a/src/BayesNet/Proposal.cc b/src/BayesNet/Proposal.cc
index eff80d4..2156b9a 100644
--- a/src/BayesNet/Proposal.cc
+++ b/src/BayesNet/Proposal.cc
@@ -12,6 +12,7 @@ namespace bayesnet {
     void Proposal::localDiscretizationProposal(map<string, vector<int>>& states, Network& model)
     {
         // order of local discretization is important. no good 0, 1, 2...
+        // although we rediscretize features after the local discretization of every feature
         auto order = model.topological_sort();
         auto& nodes = model.getNodes();
         vector<int> indicesToReDiscretize;
@@ -22,7 +23,7 @@ namespace bayesnet {
             if (nodeParents.size() < 2) continue; // Only has class as parent
             upgrade = true;
             int index = find(pFeatures.begin(), pFeatures.end(), feature) - pFeatures.begin();
-            indicesToReDiscretize.push_back(index);
+            indicesToReDiscretize.push_back(index); // We need to re-discretize this feature
             vector<string> parents;
             transform(nodeParents.begin(), nodeParents.end(), back_inserter(parents), [](const auto& p) { return p->getName(); });
             // Remove class as parent as it will be added later
@@ -30,7 +31,7 @@ namespace bayesnet {
             // Get the indices of the parents
             vector<int> indices;
             transform(parents.begin(), parents.end(), back_inserter(indices), [&](const auto& p) {return find(pFeatures.begin(), pFeatures.end(), p) - pFeatures.begin(); });
-            // Now we fit the discretizer of the feature conditioned on its parents and the class i.e. discretizer.fit(X[index], X[indices] + y)
+            // Now we fit the discretizer of the feature, conditioned on its parents and the class i.e. discretizer.fit(X[index], X[indices] + y)
             vector<string> yJoinParents;
             transform(yv.begin(), yv.end(), back_inserter(yJoinParents), [&](const auto& p) {return to_string(p); });
             for (auto idx : indices) {
@@ -43,19 +44,28 @@ namespace bayesnet {
             auto xvf_ptr = Xf.index({ index }).data_ptr<float>();
             auto xvf = vector<mdlp::precision_t>(xvf_ptr, xvf_ptr + Xf.size(1));
             discretizers[feature]->fit(xvf, yxv);
+            //
+            //
+            //
+            auto tmp = discretizers[feature]->transform(xvf);
+            Xv[index] = tmp;
+            auto xStates = vector<int>(discretizers[pFeatures[index]]->getCutPoints().size() + 1);
+            iota(xStates.begin(), xStates.end(), 0);
+            //Update new states of the feature/node
+            states[feature] = xStates;
         }
-        if (upgrade) {
-            // Discretize again X (only the affected indices) with the new fitted discretizers
-            for (auto index : indicesToReDiscretize) {
-                auto Xt_ptr = Xf.index({ index }).data_ptr<float>();
-                auto Xt = vector<float>(Xt_ptr, Xt_ptr + Xf.size(1));
-                Xv[index] = discretizers[pFeatures[index]]->transform(Xt);
-                auto xStates = vector<int>(discretizers[pFeatures[index]]->getCutPoints().size() + 1);
-                iota(xStates.begin(), xStates.end(), 0);
-                //Update new states of the feature/node
-                states[pFeatures[index]] = xStates;
-            }
-        }
+        // if (upgrade) {
+        //     // Discretize again X (only the affected indices) with the new fitted discretizers
+        //     for (auto index : indicesToReDiscretize) {
+        //         auto Xt_ptr = Xf.index({ index }).data_ptr<float>();
+        //         auto Xt = vector<float>(Xt_ptr, Xt_ptr + Xf.size(1));
+        //         Xv[index] = discretizers[pFeatures[index]]->transform(Xt);
+        //         auto xStates = vector<int>(discretizers[pFeatures[index]]->getCutPoints().size() + 1);
+        //         iota(xStates.begin(), xStates.end(), 0);
+        //         //Update new states of the feature/node
+        //         states[pFeatures[index]] = xStates;
+        //     }
+        // }
     }
     void Proposal::fit_local_discretization(map<string, vector<int>>& states, torch::Tensor& y)
     {
diff --git a/src/BayesNet/TANNew.cc b/src/BayesNet/TANNew.cc
index 44b5658..abee474 100644
--- a/src/BayesNet/TANNew.cc
+++ b/src/BayesNet/TANNew.cc
@@ -3,7 +3,6 @@
 namespace bayesnet {
     using namespace std;
     TANNew::TANNew() : TAN(), Proposal(TAN::Xv, TAN::yv, TAN::features, TAN::className) {}
-    TANNew::~TANNew() {}
     TANNew& TANNew::fit(torch::Tensor& X_, torch::Tensor& y_, vector<string>& features_, string className_, map<string, vector<int>>& states_)
     {
         // This first part should go in a Classifier method called fit_local_discretization o fit_float...
@@ -17,8 +16,9 @@ namespace bayesnet {
         cout << "TANNew: Fitting model" << endl;
         TAN::fit(TAN::Xv, TAN::yv, TAN::features, TAN::className, states);
         cout << "TANNew: Model fitted" << endl;
-        //localDiscretizationProposal(states, model);
-        //addNodes();
+        localDiscretizationProposal(states, model);
+        addNodes();
+        model.fit(TAN::Xv, TAN::yv, features, className);
         return *this;
     }
     Tensor TANNew::predict(Tensor& X)
diff --git a/src/BayesNet/TANNew.h b/src/BayesNet/TANNew.h
index 78e5f70..40609e7 100644
--- a/src/BayesNet/TANNew.h
+++ b/src/BayesNet/TANNew.h
@@ -9,7 +9,7 @@ namespace bayesnet {
     private:
     public:
         TANNew();
-        virtual ~TANNew();
+        virtual ~TANNew() = default;
         TANNew& fit(torch::Tensor& X, torch::Tensor& y, vector<string>& features, string className, map<string, vector<int>>& states) override;
         vector<string> graph(const string& name = "TAN") override;
         Tensor predict(Tensor& X) override;