fix set device.

dmlc · Feb 3, 2023 · 202c058 · 202c058
1 parent 1dbc98b
commit 202c058
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/objective/adaptive.h b/src/objective/adaptive.h
@@ -100,6 +100,7 @@ inline void UpdateTreeLeaf(Context const* ctx, HostDeviceVector<bst_node_t> cons
     detail::UpdateTreeLeafHost(ctx, position.ConstHostVector(), group_idx, info, predt, alpha,
                                p_tree);
   } else {
+    position.SetDevice(ctx->gpu_id);
     detail::UpdateTreeLeafDevice(ctx, position.ConstDeviceSpan(), group_idx, info, predt, alpha,
                                  p_tree);
   }