Align device id in predict transform with predictor. (#6662)

2021-02-02 08:33:29 +08:00
parent d8ec7aad5a
commit a9ec0ea6da
6 changed files with 30 additions and 8 deletions
--- a/src/objective/aft_obj.cu
+++ b/src/objective/aft_obj.cu
@@ -108,7 +108,7 @@ class AFTObj : public ObjFunction {
        [] XGBOOST_DEVICE(size_t _idx, common::Span<bst_float> _preds) {
      _preds[_idx] = exp(_preds[_idx]);
    }, common::Range{0, static_cast<int64_t>(io_preds->Size())},
-        tparam_->gpu_id)
+        io_preds->DeviceIdx())
    .Eval(io_preds);
  }