fixed async stream

rapidsai · Dec 13, 2019 · 8fe7326 · 8fe7326
1 parent 0e76d36
commit 8fe7326
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cpp/src/decisiontree/levelalgo/common_helper.cuh b/cpp/src/decisiontree/levelalgo/common_helper.cuh
@@ -183,8 +183,8 @@ void convert_scatter_to_gather(const unsigned int *flagsptr,
                                unsigned int *nodecount, unsigned int *nodestart,
                                unsigned int *samplelist,
                                std::shared_ptr<TemporaryMemory<T, L>> tempmem) {
-  CUDA_CHECK(
-    cudaMemsetAsync(nodestart, 0, (n_nodes + 1) * sizeof(unsigned int)));
+  CUDA_CHECK(cudaMemsetAsync(nodestart, 0, (n_nodes + 1) * sizeof(unsigned int),
+                             tempmem->stream));
   CUDA_CHECK(cudaMemsetAsync(nodecount, 0, (n_nodes + 1) * sizeof(unsigned int),
                              tempmem->stream));