update DecisionForestRegressorImpl.cpp

Signed-off-by: minmingzhu <minming.zhu@intel.com>
oap-project · Aug 9, 2023 · a88601f · a88601f
1 parent a404670
commit a88601f
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/mllib-dal/src/main/native/DecisionForestRegressorImpl.cpp b/mllib-dal/src/main/native/DecisionForestRegressorImpl.cpp
@@ -221,8 +221,8 @@ static jobject doRFRegressorOneAPICompute(
     std::cout << "doRFRegressorOneAPICompute get_column_count = "
               << hFeaturetable.get_column_count() << std::endl;
     const auto &dtype = hFeaturetable.get_metadata().get_data_type(0);
-    df::train_result result_train;
-    df::infer_result result_infer;
+    df::train_result<df::task::regression> result_train;
+    df::infer_result<df::task::regression> result_infer;
     auto t1 = std::chrono::high_resolution_clock::now();
     auto t2 = std::chrono::high_resolution_clock::now();
     auto duration =
@@ -275,7 +275,7 @@ static jobject doRFRegressorOneAPICompute(
         duration = (float)std::chrono::duration_cast<std::chrono::milliseconds>(
                        t2 - t1)
                        .count();
-        std::cout << "DF Classifier (native): training step took "
+        std::cout << "DF Regression (native): training step took "
                   << duration / 1000 << " secs." << std::endl;
         result_infer = preview::infer(comm, df_desc, result_train.get_model(),
                                       hFeaturetable);
@@ -293,7 +293,12 @@ static jobject doRFRegressorOneAPICompute(
         std::cout << "OOB error: " << result_train.get_oob_err() << std::endl;
         std::cout << "Prediction results:\n"
                   << result_infer.get_responses() << std::endl;
-
+        t2 = std::chrono::high_resolution_clock::now();
+        duration = (float)std::chrono::duration_cast<std::chrono::milliseconds>(
+                       t2 - t1)
+                       .count();
+        std::cout << "DF Regression (native): training step took "
+                  << duration / 1000 << " secs." << std::endl;
         // convert c++ map to java hashmap
         jint statsSize = 3; // spark create VarianceCalculator needs array of
                             // sufficient statistics