fix compilation error for test_gpu_predictor.cu

razdoburdin · Nov 30, 2023 · ae078a1 · ae078a1
1 parent ac5a432
commit ae078a1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/cpp/predictor/test_gpu_predictor.cu b/tests/cpp/predictor/test_gpu_predictor.cu
@@ -294,7 +294,8 @@ TEST(GPUPredictor, CategoricalPredictLeaf) {
 }
 
 TEST_F(MGPUPredictorTest, CategoricalPredictionLeafColumnSplit) {
-  RunWithInMemoryCommunicator(world_size_, TestCategoricalPredictLeaf, true, true);
+  auto ctx = MakeCUDACtx(common::AllVisibleGPUs() == 1 ? 0 : collective::GetRank());
+  RunWithInMemoryCommunicator(world_size_, TestCategoricalPredictLeaf, &ctx, true);
 }
 
 TEST(GPUPredictor, PredictLeafBasic) {