Skip zero for sparse matrix in treelite input (#297)

neo-ai · Nov 30, 2020 · 097c74b · 097c74b
1 parent f7855da
commit 097c74b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/dlr_treelite.cc b/src/dlr_treelite.cc
@@ -139,7 +139,7 @@ void TreeliteModel::SetInput(const char* name, const int64_t* shape, const void*
   treelite_input_->row_ptr.reserve(batch_size);
   for (size_t i = 0; i < batch_size; ++i) {
     for (uint32_t j = 0; j < num_col; ++j) {
-      if (!std::isnan(input_f[i * num_col + j])) {
+      if (!std::isnan(input_f[i * num_col + j]) && input_f[i * num_col + j] != 0.0f) {
         treelite_input_->data.push_back(input_f[i * num_col + j]);
         treelite_input_->col_ind.push_back(j);
       }