Fix GPU Predictor.

dmlc · Feb 3, 2020 · def3135 · def3135
1 parent fc6c123
commit def3135
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu
@@ -308,16 +308,20 @@ class GPUPredictor : public xgboost::Predictor {
                     const gbm::GBTreeModel& model, int tree_begin,
                     unsigned ntree_limit = 0) override {
     int device = generic_param_->gpu_id;
-    auto* out_preds = &predts->predictions;
     CHECK_GE(device, 0) << "Set `gpu_id' to positive value for processing GPU data.";
     ConfigureDevice(device);
 
+    auto* out_preds = &predts->predictions;
     if (predts->version == 0) {
-      CHECK_EQ(out_preds->Size(), 0);
       this->InitOutPredictions(dmat->Info(), out_preds, model);
     }
 
+    ntree_limit *= model.learner_model_param_->num_output_group;
+    if (ntree_limit == 0 || ntree_limit > model.trees.size()) {
+      ntree_limit = static_cast<unsigned>(model.trees.size());
+    }
     auto tree_end = tree_begin + ntree_limit;
+    CHECK_GE(tree_end, predts->version);
     if (tree_end - predts->version == 0) {
       CHECK_EQ(out_preds->Size(), dmat->Info().num_row_);
     } else {