Use array interface in Python prediction return. (#9855)

This commit is contained in:
Jiaming Yuan
2023-12-08 03:42:14 +08:00
committed by GitHub
parent 2c0fc97306
commit 39c637ee19
4 changed files with 49 additions and 48 deletions

View File

@@ -1097,6 +1097,8 @@ inline void CUDAEvent::Record(CUDAStreamView stream) { // NOLINT
dh::safe_cuda(cudaEventRecord(event_, cudaStream_t{stream}));
}
// Changing this has effect on prediction return, where we need to pass the pointer to
// third-party libraries like cuPy
inline CUDAStreamView DefaultStream() {
#ifdef CUDA_API_PER_THREAD_DEFAULT_STREAM
return CUDAStreamView{cudaStreamPerThread};