diff --git a/sample/CMakeLists.txt b/sample/CMakeLists.txt
index 62533c7..28f910a 100644
--- a/sample/CMakeLists.txt
+++ b/sample/CMakeLists.txt
@@ -1,4 +1,3 @@
 include_directories(${BayesNet_SOURCE_DIR}/src)
-link_directories(${MyProject_SOURCE_DIR}/src)
 add_executable(main main.cc ArffFiles.cc CPPFImdlp.cpp Metrics.cpp)
 target_link_libraries(main BayesNet "${TORCH_LIBRARIES}")
\ No newline at end of file
diff --git a/src/BaseClassifier.cc b/src/BaseClassifier.cc
index 7cc3814..21b69a1 100644
--- a/src/BaseClassifier.cc
+++ b/src/BaseClassifier.cc
@@ -8,7 +8,6 @@ namespace bayesnet {
     BaseClassifier::BaseClassifier(Network model) : model(model), m(0), n(0), metrics(Metrics()), fitted(false) {}
     BaseClassifier& BaseClassifier::build(vector<string>& features, string className, map<string, vector<int>>& states)
     {
-
         dataset = torch::cat({ X, y.view({y.size(0), 1}) }, 1);
         this->features = features;
         this->className = className;
@@ -116,4 +115,13 @@ namespace bayesnet {
         }
         model.addNode(className, states[className].size());
     }
+    int BaseClassifier::getNumberOfNodes()
+    {
+        // Features does not include class
+        return fitted ? model.getFeatures().size() + 1 : 0;
+    }
+    int BaseClassifier::getNumberOfEdges()
+    {
+        return fitted ? model.getEdges().size() : 0;
+    }
 }
\ No newline at end of file
diff --git a/src/BaseClassifier.h b/src/BaseClassifier.h
index 730f3cd..dcb0b34 100644
--- a/src/BaseClassifier.h
+++ b/src/BaseClassifier.h
@@ -30,6 +30,8 @@ namespace bayesnet {
         virtual ~BaseClassifier() = default;
         BaseClassifier& fit(vector<vector<int>>& X, vector<int>& y, vector<string>& features, string className, map<string, vector<int>>& states);
         void addNodes();
+        int getNumberOfNodes();
+        int getNumberOfEdges();
         Tensor predict(Tensor& X);
         vector<int> predict(vector<vector<int>>& X);
         float score(Tensor& X, Tensor& y);
diff --git a/src/Network.cc b/src/Network.cc
index b6cda9c..a4f678d 100644
--- a/src/Network.cc
+++ b/src/Network.cc
@@ -275,5 +275,17 @@ namespace bayesnet {
         output.push_back("}\n");
         return output;
     }
+    vector<pair<string, string>> Network::getEdges()
+    {
+        auto edges = vector<pair<string, string>>();
+        for (const auto& node : nodes) {
+            auto head = node.first;
+            for (const auto& child : node.second->getChildren()) {
+                auto tail = child->getName();
+                edges.push_back({ head, tail });
+            }
+        }
+        return edges;
+    }
 
 }
diff --git a/src/Network.h b/src/Network.h
index e3b3b19..e7e8b68 100644
--- a/src/Network.h
+++ b/src/Network.h
@@ -36,6 +36,7 @@ namespace bayesnet {
         map<string, std::unique_ptr<Node>>& getNodes();
         vector<string> getFeatures();
         int getStates();
+        vector<pair<string, string>> getEdges();
         int getClassNumStates();
         string getClassName();
         void fit(const vector<vector<int>>&, const vector<int>&, const vector<string>&, const string&);
diff --git a/tests/main.cc b/tests/BayesModels.cc
similarity index 64%
rename from tests/main.cc
rename to tests/BayesModels.cc
index 4dfe3d2..1d20edb 100644
--- a/tests/main.cc
+++ b/tests/BayesModels.cc
@@ -21,7 +21,7 @@ TEST_CASE("Test Bayesian Classifiers score", "[BayesNet]")
     };
 
     string file_name = GENERATE("glass", "iris", "ecoli", "diabetes");
-    auto[Xd, y, features, className, states] = loadFile(file_name);
+    auto [Xd, y, features, className, states] = loadFile(file_name);
 
     SECTION("Test TAN classifier (" + file_name + ")")
     {
@@ -59,4 +59,30 @@ TEST_CASE("Test Bayesian Classifiers score", "[BayesNet]")
     // for (auto scores : scores) {
     //     cout << "{{\"" << scores.first.first << "\", \"" << scores.first.second << "\"}, " << scores.second << "}, ";
     // }
+}
+TEST_CASE("Models features")
+{
+    auto graph = vector<string>({ "digraph BayesNet {\nlabel=<BayesNet Test>\nfontsize=30\nfontcolor=blue\nlabelloc=t\nlayout=circo\n",
+        "class [shape=circle, fontcolor=red, fillcolor=lightblue, style=filled ] \n",
+        "class -> sepallength", "class -> sepalwidth", "class -> petallength", "class -> petalwidth", "petallength [shape=circle] \n",
+        "petallength -> sepallength", "petalwidth [shape=circle] \n", "sepallength [shape=circle] \n",
+        "sepallength -> sepalwidth", "sepalwidth [shape=circle] \n", "sepalwidth -> petalwidth", "}\n"
+        }
+    );
+
+    auto clf = bayesnet::TAN();
+    auto [Xd, y, features, className, states] = loadFile("iris");
+    clf.fit(Xd, y, features, className, states);
+    REQUIRE(clf.getNumberOfNodes() == 5);
+    REQUIRE(clf.getNumberOfEdges() == 7);
+    REQUIRE(clf.show() == vector<string>{"class -> sepallength, sepalwidth, petallength, petalwidth, ", "petallength -> sepallength, ", "petalwidth -> ", "sepallength -> sepalwidth, ", "sepalwidth -> petalwidth, "});
+    REQUIRE(clf.graph("Test") == graph);
+}
+TEST_CASE("Get num features & num edges")
+{
+    auto [Xd, y, features, className, states] = loadFile("iris");
+    auto clf = bayesnet::KDB(2);
+    clf.fit(Xd, y, features, className, states);
+    REQUIRE(clf.getNumberOfNodes() == 5);
+    REQUIRE(clf.getNumberOfEdges() == 8);
 }
\ No newline at end of file
diff --git a/tests/BayesNetwork.cc b/tests/BayesNetwork.cc
index b6bda60..f9052b7 100644
--- a/tests/BayesNetwork.cc
+++ b/tests/BayesNetwork.cc
@@ -7,7 +7,7 @@
 
 TEST_CASE("Test Bayesian Network")
 {
-    auto[Xd, y, features, className, states] = loadFile("iris");
+    auto [Xd, y, features, className, states] = loadFile("iris");
 
     SECTION("Test Update Nodes")
     {
@@ -26,4 +26,16 @@ TEST_CASE("Test Bayesian Network")
         net.addNode("C", 2);
         REQUIRE(net.getFeatures() == vector<string>{"A", "B", "C"});
     }
+    SECTION("Test get edges")
+    {
+        auto net = bayesnet::Network();
+        net.addNode("A", 3);
+        net.addNode("B", 5);
+        net.addNode("C", 2);
+        net.addEdge("A", "B");
+        net.addEdge("B", "C");
+        REQUIRE(net.getEdges() == vector<pair<string, string>>{ {"A", "B"}, { "B", "C" } });
+        net.addEdge("A", "C");
+        REQUIRE(net.getEdges() == vector<pair<string, string>>{ {"A", "B"}, { "A", "C" }, { "B", "C" } });
+    }
 }
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index e7f0438..3278d83 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -1,8 +1,8 @@
 if(ENABLE_TESTING)
     set(TEST_MAIN "unit_tests")
-    set(TEST_SOURCES main.cc ../sample/ArffFiles.cc ../sample/CPPFImdlp.cpp ../sample/Metrics.cpp 
+    set(TEST_SOURCES BayesModels.cc BayesNetwork.cc ../sample/ArffFiles.cc ../sample/CPPFImdlp.cpp ../sample/Metrics.cpp 
     ../src/utils.cc ../src/Network.cc ../src/Node.cc ../src/Metrics.cc ../src/BaseClassifier.cc ../src/KDB.cc 
-    ../src/TAN.cc ../src/SPODE.cc ../src/Ensemble.cc ../src/AODE.cc ../src/Mst.cc BayesNetwork.cc utils.cc utils.h)
+    ../src/TAN.cc ../src/SPODE.cc ../src/Ensemble.cc ../src/AODE.cc ../src/Mst.cc utils.cc utils.h)
     add_executable(${TEST_MAIN} ${TEST_SOURCES})
     target_link_libraries(${TEST_MAIN} PUBLIC "${TORCH_LIBRARIES}"  Catch2::Catch2WithMain)
     add_test(NAME ${TEST_MAIN} COMMAND ${TEST_MAIN})