Allocate a temp buffer for binary predict_proba

Author: hcho3

src/cpu_forest_model.h

@@ -82,20 +82,21 @@ struct ForestModel<rapids::HostMemory> {
       auto output_buffer = rapids::Buffer<float>{
           output.data(), output.size(), output.mem_type(), output.device(),
           output.stream()};
-      auto output_size = output.size();
+      auto const num_classes = tl_model_->num_classes();
       // New FIL expects buffer of size samples * num_classes for multi-class
       // classifiers, but output buffer may be smaller, so we need a temporary
       // buffer
-      auto const num_classes = tl_model_->num_classes();
       if (!predict_proba && tl_model_->config().output_class &&
           num_classes > 1) {
-        output_size = samples * num_classes;
-        if (output_size != output.size()) {
-          // If expected output size is not the same as the size of `output`,
-          // create a temporary buffer of the correct size
-          output_buffer =
-              rapids::Buffer<float>{output_size, rapids::HostMemory};
-        }
+        output_buffer =
+            rapids::Buffer<float>{samples * num_classes, rapids::HostMemory};
+      } else if (
+          predict_proba && tl_model_->config().output_class &&
+          num_classes == 1) {
+        // Also use a temp buffer when probabilities are requested for
+        // a binary classifier. This is so that we can output probabilities
+        // for both positive and negative classes.
+        output_buffer = rapids::Buffer<float>{samples * 2, rapids::HostMemory};
       }
 
       // TODO(hcho3): Revise new FIL so that it takes in (const io_t*) type for

src/gpu_forest_model.h

@@ -108,20 +108,22 @@ struct ForestModel<rapids::DeviceMemory> {
       auto output_buffer = rapids::Buffer<float>{
           output.data(), output.size(), output.mem_type(), output.device(),
           output.stream()};
-      auto output_size = output.size();
+      auto const num_classes = tl_model_->num_classes();
       // New FIL expects buffer of size samples * num_classes for multi-class
       // classifiers, but output buffer may be smaller, so we need a temporary
       // buffer
-      auto const num_classes = tl_model_->num_classes();
       if (!predict_proba && tl_model_->config().output_class &&
           num_classes > 1) {
-        output_size = samples * num_classes;
-        if (output_size != output.size()) {
-          // If expected output size is not the same as the size of `output`,
-          // create a temporary buffer of the correct size
-          output_buffer =
-              rapids::Buffer<float>{output_size, rapids::DeviceMemory};
-        }
+        output_buffer =
+            rapids::Buffer<float>{samples * num_classes, rapids::DeviceMemory};
+      } else if (
+          predict_proba && tl_model_->config().output_class &&
+          num_classes == 1) {
+        // Also use a temp buffer when probabilities are requested for
+        // a binary classifier. This is so that we can output probabilities
+        // for both positive and negative classes.
+        output_buffer =
+            rapids::Buffer<float>{samples * 2, rapids::DeviceMemory};
       }
       // TODO(hcho3): Revise new FIL so that it takes in (const io_t*) type for
       // input buffer