Specify the allocator.

dmlc · Jun 23, 2021 · 3050cfa · 3050cfa
1 parent 957d778
commit 3050cfa
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/gputreeshap b/gputreeshap
diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu
@@ -788,8 +788,9 @@ class GPUPredictor : public xgboost::Predictor {
       SparsePageView X(batch.data.DeviceSpan(), batch.offset.DeviceSpan(),
                        model.learner_model_param->num_feature);
       auto begin = dh::tbegin(phis) + batch.base_rowid * contributions_columns;
-      gpu_treeshap::GPUTreeShap(X, device_paths.begin(), device_paths.end(),
-                                ngroup, begin, dh::tend(phis));
+      gpu_treeshap::GPUTreeShap<dh::XGBDeviceAllocator<int>>(
+          X, device_paths.begin(), device_paths.end(), ngroup, begin,
+          dh::tend(phis));
     }
     // Add the base margin term to last column
     p_fmat->Info().base_margin_.SetDevice(generic_param_->gpu_id);
@@ -844,9 +845,9 @@ class GPUPredictor : public xgboost::Predictor {
       SparsePageView X(batch.data.DeviceSpan(), batch.offset.DeviceSpan(),
                        model.learner_model_param->num_feature);
       auto begin = dh::tbegin(phis) + batch.base_rowid * contributions_columns;
-      gpu_treeshap::GPUTreeShapInteractions(X, device_paths.begin(),
-                                            device_paths.end(), ngroup, begin,
-                                            dh::tend(phis));
+      gpu_treeshap::GPUTreeShapInteractions<dh::XGBDeviceAllocator<int>>(
+          X, device_paths.begin(), device_paths.end(), ngroup, begin,
+          dh::tend(phis));
     }
     // Add the base margin term to last column
     p_fmat->Info().base_margin_.SetDevice(generic_param_->gpu_id);