diff --git a/gputreeshap b/gputreeshap
index b6ecc9e57cea..d25aa1be9630 160000
--- a/gputreeshap
+++ b/gputreeshap
@@ -1 +1 @@
-Subproject commit b6ecc9e57cea8d638cba8ceeb70261b9a0c0eb6e
+Subproject commit d25aa1be9630f423c029ec20fe95557698dc7e42
diff --git a/src/predictor/gpu_predictor.cu b/src/predictor/gpu_predictor.cu
index 52456a62272f..64a4ffd6de2f 100644
--- 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);