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

tests/cpp/objective/test_aft_obj.cc

@@ -49,7 +49,7 @@ static inline void CheckGPairOverGridPoints(
   const int num_point = 20;
   const double log_y_low = 1.0;
   const double log_y_high = 15.0;
-  
+
   obj->Configure({ {"aft_loss_distribution", dist_type},
                    {"aft_loss_distribution_scale", "1"} });
 

tests/cpp/objective/test_objective.cc

@@ -17,3 +17,25 @@ TEST(Objective, UnknownFunction) {
     delete obj;
   }
 }
+
+namespace xgboost {
+TEST(Objective, PredTransform) {
+  // Test that show PredTransform uses the same device with predictor.
+  xgboost::GenericParameter tparam;
+  tparam.UpdateAllowUnknown(Args{{"gpu_id", "0"}});
+  size_t n = 100;
+
+  for (const auto &entry :
+       ::dmlc::Registry<::xgboost::ObjFunctionReg>::List()) {
+    std::unique_ptr<xgboost::ObjFunction> obj{
+        xgboost::ObjFunction::Create(entry->name, &tparam)};
+    obj->Configure(Args{{"num_class", "2"}});
+    HostDeviceVector<float> predts;
+    predts.Resize(n, 3.14f);  // prediction is performed on host.
+    ASSERT_FALSE(predts.DeviceCanRead());
+    obj->PredTransform(&predts);
+    ASSERT_FALSE(predts.DeviceCanRead());
+    ASSERT_TRUE(predts.HostCanWrite());
+  }
+}
+} // namespace xgboost