1 files changed, 9 insertions, 8 deletions
diff --git a/main.cpp b/main.cpp
index c297f30..b479f59 100644
--- a/main.cpp
+++ b/main.cpp
@@ -5,6 +5,7 @@
 #include "tools/easylogging++.h"
 #include "model/ranksvmtn.h"
 #include "tools/fileDataProvider.h"
+#include "tools/matrixIO.h"
 
 INITIALIZE_EASYLOGGINGPP
 
@@ -22,17 +23,15 @@ int train() {
     dp.open();
     DataSet D;
     Labels L;
+
     LOG(INFO)<<"Training started";
-    while (!dp.EOFile())
-    {
-        dp.getDataSet(D);
-        dp.getLabel(L);
-        rsvm->train(D,L);
-    }
+    dp.getDataSet(D);
+    dp.getLabel(L);
+    rsvm->train(D,L);
 
     LOG(INFO)<<"Training finished,saving model";
 
-
+    dp.close();
     rsvm->saveModel(vm["output"].as<std::string>().c_str());
     delete rsvm;
     return 0;
@@ -43,12 +42,14 @@ int predict() {
     rsvm = RSVM::loadModel(vm["model"].as<std::string>().c_str());
     FileDP dp(vm["feature"].as<std::string>().c_str());
     DataSet D;
-    MatrixXd L;
+    Labels L;
     while (!dp.EOFile())
     {
         dp.getDataSet(D);
         rsvm->predict(D,L);
     }
+
+    Eigen::write_stream(std::cout, L);
     delete rsvm;
     return 0;
 }