diff --git a/src/PyClassifier.cc b/src/PyClassifier.cc
index 39b8885..c7684a4 100644
--- a/src/PyClassifier.cc
+++ b/src/PyClassifier.cc
@@ -3,7 +3,7 @@
 #include <iostream>
 
 namespace pywrap {
-    namespace p = boost::python;
+    namespace bp = boost::python;
     namespace np = boost::python::numpy;
     PyClassifier::PyClassifier(const std::string& module, const std::string& className) : module(module), className(className)
     {
@@ -20,14 +20,14 @@ namespace pywrap {
     {
         int m = X.size(0);
         int n = X.size(1);
-        auto Xn = np::from_data(X.data_ptr(), np::dtype::get_builtin<float>(), p::make_tuple(m, n), p::make_tuple(sizeof(X.dtype()) * 2 * n, sizeof(X.dtype()) * 2), p::object());
+        auto Xn = np::from_data(X.data_ptr(), np::dtype::get_builtin<float>(), bp::make_tuple(m, n), bp::make_tuple(sizeof(X.dtype()) * 2 * n, sizeof(X.dtype()) * 2), bp::object());
         Xn = Xn.transpose();
         return Xn;
     }
     std::pair<np::ndarray, np::ndarray> tensors2numpy(torch::Tensor& X, torch::Tensor& y)
     {
         int n = X.size(1);
-        auto yn = np::from_data(y.data_ptr(), np::dtype::get_builtin<int32_t>(), p::make_tuple(n), p::make_tuple(sizeof(y.dtype()) * 2), p::object());
+        auto yn = np::from_data(y.data_ptr(), np::dtype::get_builtin<int32_t>(), bp::make_tuple(n), bp::make_tuple(sizeof(y.dtype()) * 2), bp::object());
         return { tensor2numpy(X), yn };
     }
     std::string PyClassifier::version()
@@ -41,50 +41,36 @@ namespace pywrap {
     PyClassifier& PyClassifier::fit(torch::Tensor& X, torch::Tensor& y, const std::vector<std::string>& features, const std::string& className, std::map<std::string, std::vector<int>>& states)
     {
         auto [Xn, yn] = tensors2numpy(X, y);
-        CPyObject Xp = p::incref(p::object(Xn).ptr());
-        CPyObject yp = p::incref(p::object(yn).ptr());
+        CPyObject Xp = bp::incref(bp::object(Xn).ptr());
+        CPyObject yp = bp::incref(bp::object(yn).ptr());
         pyWrap->fit(module, this->className, Xp, yp);
         return *this;
     }
-    void print_array(np::ndarray& array)
-    {
-        std::cout << "Array: " << std::endl;
-        std::cout << p::extract<char const*>(p::str(array)) << std::endl;
-    }
     torch::Tensor PyClassifier::predict(torch::Tensor& X)
     {
         int dimension = X.size(1);
         auto Xn = tensor2numpy(X);
-        CPyObject Xp = p::incref(p::object(Xn).ptr());
+        CPyObject Xp = bp::incref(bp::object(Xn).ptr());
         PyObject* incoming = pyWrap->predict(module, className, Xp);
         std::cout << "Return from predict" << std::endl;
-        p::handle<> handle(incoming);
-        p::object object(handle);
+        bp::handle<> handle(incoming);
+        bp::object object(handle);
         np::ndarray prediction = np::from_object(object);
-        print_array(prediction);
-        // import_array();
-        // if (!PyArray_Check(incoming)) {
-        //     throw std::logic_error("Returned value is not array");
-        // }
-        // std::cout << "Returned value is array" << std::endl;
-        // PyArrayObject* np_ret = (PyArrayObject*)incoming;
-        // if (PyArray_NDIM(np_ret) != dimension - 1) {
-        //     throw std::logic_error("Returned array has wrong dimension" + std::to_string(PyArray_NDIM(np_ret)) + "!=" + std::to_string(dimension - 1));
-        // }
-        // std::cout << "Returned array has correct dimension" << PyArray_NDIM(np_ret) << std::endl;
-        // int len{ PyArray_SHAPE(np_ret)[0] };
-        // int* data = reinterpret_cast<int*>(PyArray_DATA(np_ret));
-
-        // int* data = reinterpret_cast<int*>(prediction.get_data());
-        // auto resultTensor = torch::tensor({ data }, torch::kInt32);
-        auto resultTensor = torch::zeros({ prediction.shape(0) }, torch::kInt32);
+        if (PyErr_Occurred()) {
+            PyErr_Print();
+            throw std::runtime_error("Error cleaning module " + module + " and class " + className);
+        }
+        int* data = reinterpret_cast<int*>(prediction.get_data());
+        std::vector<int> v1(data, data + prediction.shape(0));
+        auto resultTensor = torch::tensor(v1, torch::kInt32);
+        Py_XDECREF(incoming);
         return resultTensor;
     }
     double PyClassifier::score(torch::Tensor& X, torch::Tensor& y)
     {
         auto [Xn, yn] = tensors2numpy(X, y);
-        CPyObject Xp = p::incref(p::object(Xn).ptr());
-        CPyObject yp = p::incref(p::object(yn).ptr());
+        CPyObject Xp = bp::incref(bp::object(Xn).ptr());
+        CPyObject yp = bp::incref(bp::object(yn).ptr());
         auto result = pyWrap->score(module, className, Xp, yp);
         return result;
     }
diff --git a/src/main.cc b/src/main.cc
index 736b194..56b0015 100644
--- a/src/main.cc
+++ b/src/main.cc
@@ -58,7 +58,11 @@ int main(int argc, char* argv[])
         clf.fit(X, y, features, className, states);
         // cout << "STree Score: " << clf.score(X, y) << endl;
         auto prediction = clf.predict(X);
-        cout << "Prediction: " << prediction << endl;
+        cout << "Prediction: " << endl << "{";
+        for (int i = 0; i < prediction.size(0); ++i) {
+            cout << prediction[i].item<int>() << ", ";
+        }
+        cout << "}" << endl;
     }
     cout << "* End." << endl;
 }
\ No newline at end of file