Check.

dmlc · Mar 24, 2021 · f3fe1d9 · f3fe1d9
1 parent 80f79dc
commit f3fe1d9
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/data/array_interface.cu b/src/data/array_interface.cu
@@ -6,6 +6,16 @@
 
 namespace xgboost {
 void ArrayInterfaceHandler::SyncCudaStream(int64_t stream) {
-  dh::safe_cuda(cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)));
+  switch (stream) {
+  case 0:
+    LOG(FATAL) << "Invalid stream ID in array interface: " << stream;
+  case 1:
+    // default legacy stream
+    break;
+  case 2:
+    // default per-thread stream
+  default:
+    dh::safe_cuda(cudaStreamSynchronize(reinterpret_cast<cudaStream_t>(stream)));
+  }
 }
 }  // namespace xgboost