From 7bf294a19a712fb36d9f1c89b1587b11c5237733 Mon Sep 17 00:00:00 2001
From: jinfeng <jinfengl@nvidia.com>
Date: Wed, 7 Feb 2024 23:32:17 -0800
Subject: [PATCH] fix device to host copy not sync stream

---
 cpp/src/glm/qn_mg.cu | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cpp/src/glm/qn_mg.cu b/cpp/src/glm/qn_mg.cu
index 1a128100c8..c790519244 100644
--- a/cpp/src/glm/qn_mg.cu
+++ b/cpp/src/glm/qn_mg.cu
@@ -65,6 +65,7 @@ std::vector<T> distinct_mg(const raft::handle_t& handle, T* y, size_t n)
 
   std::vector<size_t> recv_counts_host(n_ranks);
   raft::copy(recv_counts_host.data(), recv_counts.data(), n_ranks, stream);
+  raft::resource::sync_stream(handle);
 
   std::vector<size_t> displs(n_ranks);
   size_t pos = 0;
@@ -84,6 +85,7 @@ std::vector<T> distinct_mg(const raft::handle_t& handle, T* y, size_t n)
 
   std::vector<T> global_unique_y_host(global_unique_y.size());
   raft::copy(global_unique_y_host.data(), global_unique_y.data(), global_unique_y.size(), stream);
+  raft::resource::sync_stream(handle);
 
   return global_unique_y_host;
 }