diff --git a/src/Network.cc b/src/Network.cc
index 458eadb..391b8af 100644
--- a/src/Network.cc
+++ b/src/Network.cc
@@ -151,13 +151,17 @@ namespace bayesnet {
             for (int col = 0; col < samples.size(); ++col) {
                 sample.push_back(samples[col][row]);
             }
-            predictions.push_back(predict_sample(sample).first);
+            vector<double> classProbabilities = predict_sample(sample);
+            // Find the class with the maximum posterior probability
+            auto maxElem = max_element(classProbabilities.begin(), classProbabilities.end());
+            int predictedClass = distance(classProbabilities.begin(), maxElem);
+            predictions.push_back(predictedClass);
         }
         return predictions;
     }
-    vector<pair<int, double>> Network::predict_proba(const vector<vector<int>>& samples)
+    vector<vector<double>> Network::predict_proba(const vector<vector<int>>& samples)
     {
-        vector<pair<int, double>> predictions;
+        vector<vector<double>> predictions;
         vector<int> sample;
         for (int row = 0; row < samples[0].size(); ++row) {
             sample.clear();
@@ -179,7 +183,7 @@ namespace bayesnet {
         }
         return (double)correct / y_pred.size();
     }
-    pair<int, double> Network::predict_sample(const vector<int>& sample)
+    vector<double> Network::predict_sample(const vector<int>& sample)
     {
         // Ensure the sample size is equal to the number of features
         if (sample.size() != features.size()) {
@@ -190,14 +194,8 @@ namespace bayesnet {
         for (int i = 0; i < sample.size(); ++i) {
             evidence[features[i]] = sample[i];
         }
-        vector<double> classProbabilities = exactInference(evidence);
+        return exactInference(evidence);
 
-        // Find the class with the maximum posterior probability
-        auto maxElem = max_element(classProbabilities.begin(), classProbabilities.end());
-        int predictedClass = distance(classProbabilities.begin(), maxElem);
-        double maxProbability = *maxElem;
-
-        return make_pair(predictedClass, maxProbability);
     }
     double Network::computeFactor(map<string, int>& completeEvidence)
     {
diff --git a/src/Network.h b/src/Network.h
index db4a1e3..78ce3ab 100644
--- a/src/Network.h
+++ b/src/Network.h
@@ -16,7 +16,7 @@ namespace bayesnet {
         string className;
         int laplaceSmoothing;
         bool isCyclic(const std::string&, std::unordered_set<std::string>&, std::unordered_set<std::string>&);
-        pair<int, double> predict_sample(const vector<int>&);
+        vector<double> predict_sample(const vector<int>&);
         vector<double> exactInference(map<string, int>&);
         double computeFactor(map<string, int>&);
     public:
@@ -34,7 +34,7 @@ namespace bayesnet {
         string getClassName();
         void fit(const vector<vector<int>>&, const vector<int>&, const vector<string>&, const string&);
         vector<int> predict(const vector<vector<int>>&);
-        vector<pair<int, double>> predict_proba(const vector<vector<int>>&);
+        vector<vector<double>> predict_proba(const vector<vector<int>>&);
         double score(const vector<vector<int>>&, const vector<int>&);
         inline string version() { return "0.1.0"; }
     };