From 24e3153bfdf522e31496374b7b870f6f8f3d5ace Mon Sep 17 00:00:00 2001
From: Hyunsu Cho <phcho@nvidia.com>
Date: Mon, 11 Sep 2023 11:27:26 -0700
Subject: [PATCH] Pass correct nthread to ClassEncoder

---
 src/cpu_forest_model.h | 5 ++++-
 src/gpu_forest_model.h | 2 ++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/cpu_forest_model.h b/src/cpu_forest_model.h
index 3a40e81..f9e403e 100644
--- a/src/cpu_forest_model.h
+++ b/src/cpu_forest_model.h
@@ -29,6 +29,7 @@
 #include <cuml/experimental/fil/detail/raft_proto/handle.hpp>
 #include <cuml/experimental/fil/forest_model.hpp>
 #include <cuml/experimental/fil/treelite_importer.hpp>
+#include <herring/omp_helpers.hpp>
 #include <memory>
 #include <optional>
 #include <rapids_triton/memory/buffer.hpp>
@@ -54,6 +55,8 @@ struct ForestModel<rapids::HostMemory> {
                   *tl_model_->base_tl_model(), filex::preferred_tree_layout,
                   filex::index_type{}, std::nullopt,
                   raft_proto::device_type::cpu);
+              rapids::log_info(__FILE__, __LINE__)
+                  << "Loaded model to new FIL format";
             }
             catch (filex::model_import_error const& ex) {
               result = std::nullopt;
@@ -65,7 +68,7 @@ struct ForestModel<rapids::HostMemory> {
           }
           return result;
         }()},
-        class_encoder_{tl_model_->config().cpu_nthread}
+        class_encoder_{int(thread_count(tl_model_->config().cpu_nthread))}
   {
   }
 
diff --git a/src/gpu_forest_model.h b/src/gpu_forest_model.h
index 0ce478f..94d47dc 100644
--- a/src/gpu_forest_model.h
+++ b/src/gpu_forest_model.h
@@ -75,6 +75,8 @@ struct ForestModel<rapids::DeviceMemory> {
                   *tl_model_->base_tl_model(), filex::preferred_tree_layout,
                   filex::index_type{}, std::nullopt,
                   raft_proto::device_type::gpu);
+              rapids::log_info(__FILE__, __LINE__)
+                  << "Loaded model to new FIL format";
             }
             catch (filex::model_import_error const& ex) {
               result = std::nullopt;