diff --git a/.vscode/launch.json b/.vscode/launch.json
index c8b7f7f..ff1e951 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -13,11 +13,9 @@
                 "TAN",
                 "-p",
                 "../../data/",
-                "--stratified",
                 "--tensors"
             ],
             "cwd": "${workspaceFolder}/build/sample/",
-            "preLaunchTask": "CMake: build"
         },
         {
             "type": "lldb",
diff --git a/sample/sample.cc b/sample/sample.cc
index 613a2fa..039a1e9 100644
--- a/sample/sample.cc
+++ b/sample/sample.cc
@@ -228,8 +228,7 @@ int main(int argc, char** argv)
         total_score += score_test;
         cout << "Score Train: " << score_train << endl;
         cout << "Score Test : " << score_test << endl;
-        // cout << "-------------------------------------------------------------------------------" << endl;
-        // total_score += score_value;
+        cout << "-------------------------------------------------------------------------------" << endl;
     }
     cout << "**********************************************************************************" << endl;
     cout << "Average Score Train: " << total_score_train / nFolds << endl;
diff --git a/src/BayesNet/Classifier.cc b/src/BayesNet/Classifier.cc
index 6fb67d9..f6297a2 100644
--- a/src/BayesNet/Classifier.cc
+++ b/src/BayesNet/Classifier.cc
@@ -7,6 +7,11 @@ namespace bayesnet {
     Classifier::Classifier(Network model) : model(model), m(0), n(0), metrics(Metrics()), fitted(false) {}
     Classifier& Classifier::build(vector<string>& features, string className, map<string, vector<int>>& states)
     {
+        cout << "Building classifier..." << endl;
+        cout << "X sizes = " << X.sizes() << endl;
+        cout << "y sizes = " << y.sizes() << endl;
+        cout << "Xv size = " << Xv.size() << endl;
+        cout << "yv size = " << yv.size() << endl;
         dataset = torch::cat({ X, y.view({y.size(0), 1}) }, 1);
         this->features = features;
         this->className = className;
@@ -28,6 +33,10 @@ namespace bayesnet {
         this->X = torch::transpose(X, 0, 1);
         this->y = y;
         Xv = vector<vector<int>>();
+        for (int i = 0; i < X.size(1); ++i) {
+            auto temp = X.index({ "...", i });
+            Xv.push_back(vector<int>(temp.data_ptr<int>(), temp.data_ptr<int>() + temp.numel()));
+        }
         yv = vector<int>(y.data_ptr<int>(), y.data_ptr<int>() + y.size(0));
         return build(features, className, states);
     }
@@ -71,10 +80,11 @@ namespace bayesnet {
         }
         auto m_ = X.size(0);
         auto n_ = X.size(1);
+        //auto Xt = torch::transpose(X, 0, 1);
         vector<vector<int>> Xd(n_, vector<int>(m_, 0));
         for (auto i = 0; i < n_; i++) {
             auto temp = X.index({ "...", i });
-            Xd[i] = vector<int>(temp.data_ptr<int>(), temp.data_ptr<int>() + m_);
+            Xd[i] = vector<int>(temp.data_ptr<int>(), temp.data_ptr<int>() + temp.numel());
         }
         auto yp = model.predict(Xd);
         auto ypred = torch::tensor(yp, torch::kInt32);
diff --git a/src/Platform/Experiment.cc b/src/Platform/Experiment.cc
index 6bca89b..464b89c 100644
--- a/src/Platform/Experiment.cc
+++ b/src/Platform/Experiment.cc
@@ -145,7 +145,6 @@ int main(int argc, char** argv)
         fold = new StratifiedKFold(n_folds, y, -1);
     else
         fold = new KFold(n_folds, y.numel(), -1);
-
     auto experiment = Experiment();
     experiment.setDiscretized(discretize_dataset).setModel(model_name).setPlatform("cpp");
     experiment.setStratified(stratified).setNFolds(5).addRandomSeed(271).setScoreName("accuracy");
diff --git a/src/Platform/Folding.cc b/src/Platform/Folding.cc
index b688353..86aaaec 100644
--- a/src/Platform/Folding.cc
+++ b/src/Platform/Folding.cc
@@ -48,7 +48,6 @@ void StratifiedKFold::build()
 {
     stratified_indices = vector<vector<int>>(k);
     int fold_size = n / k;
-    int remainder = n % k;
     // Compute class counts and indices
     auto class_indices = map<int, vector<int>>();
     vector<int> class_counts(*max_element(y.begin(), y.end()) + 1, 0);