Fix GPU inference

Author: hcho3

src/cpu_forest_model.h

@@ -100,16 +100,15 @@ struct ForestModel<rapids::HostMemory> {
         // 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
-      // input buffer
       if (new_fil_model_->row_postprocessing() == filex::row_op::max_index) {
         // Work around the limitation of max_index, to allow for probability
         // output
         // TODO(hcho3): Review new FIL to expose predict_proba(), to always
         // output probabilities
         new_fil_model_->set_row_postprocessing(filex::row_op::softmax);
       }
+      // TODO(hcho3): Revise new FIL so that it takes in (const io_t*) type for
+      // input buffer
       new_fil_model_->predict(
           raft_proto::handle_t{}, output_buffer.data(),
           const_cast<float*>(input.data()), samples,

src/gpu_forest_model.h

@@ -127,6 +127,13 @@ struct ForestModel<rapids::DeviceMemory> {
         output_buffer =
             rapids::Buffer<float>{samples * 2, rapids::DeviceMemory};
       }
+      if (new_fil_model_->row_postprocessing() == filex::row_op::max_index) {
+        // Work around the limitation of max_index, to allow for probability
+        // output
+        // TODO(hcho3): Review new FIL to expose predict_proba(), to always
+        // output probabilities
+        new_fil_model_->set_row_postprocessing(filex::row_op::softmax);
+      }
       // TODO(hcho3): Revise new FIL so that it takes in (const io_t*) type for
       // input buffer
       new_fil_model_->predict(