Optionaly fail when gpu_id is set to invalid value (#6342)

2020-11-28 08:14:12 +01:00
parent 38c80bcec4
commit 028ec5f028
4 changed files with 25 additions and 1 deletions
--- a/src/predictor/gpu_predictor.cu
+++ b/src/predictor/gpu_predictor.cu
@@ -580,7 +580,7 @@ class GPUPredictor : public xgboost::Predictor {
      Predictor::Predictor{generic_param} {}

  ~GPUPredictor() override {
-    if (generic_param_->gpu_id >= 0) {
+    if (generic_param_->gpu_id >= 0 && generic_param_->gpu_id < common::AllVisibleGPUs()) {
      dh::safe_cuda(cudaSetDevice(generic_param_->gpu_id));
    }
  }