fix bug on scikit max_mem_size

Xtra-Computing · Jul 14, 2018 · 07c812e · 07c812e
1 parent 9fae827
commit 07c812e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/thundersvm/thundersvm-scikit.cpp b/src/thundersvm/thundersvm-scikit.cpp
@@ -105,7 +105,7 @@ extern "C" {
         param_cmd.epsilon = (float_type)tol;
         param_cmd.probability = probability;
         if(max_mem_size != -1)
-            param_cmd.max_mem_size = max_mem_size;
+            param_cmd.max_mem_size = static_cast<size_t>(max(max_mem_size, 0)) << 20;
         if(weight_size != 0) {
             param_cmd.nr_weight = weight_size;
             param_cmd.weight = (float_type *) malloc(weight_size * sizeof(float_type));
@@ -201,7 +201,7 @@ extern "C" {
         param_cmd.epsilon = (float_type)tol;
         param_cmd.probability = probability;
         if(max_mem_size != -1)
-            param_cmd.max_mem_size = max_mem_size;
+            param_cmd.max_mem_size = static_cast<size_t>(max(max_mem_size, 0)) << 20;
         if(weight_size != 0) {
             param_cmd.nr_weight = weight_size;
             param_cmd.weight = (float_type *) malloc(weight_size * sizeof(float_type));