Merge pull request #70 from rapidsai/branch-0.33

Forward-merge branch-0.33 to branch-0.34
rapidsai · Jul 20, 2023 · 42e5aeb · 42e5aeb
2 parents e135aeb + c91d9e6
commit 42e5aeb
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/cpp/tests/request.cpp b/cpp/tests/request.cpp
@@ -122,6 +122,9 @@ class RequestTest : public ::testing::TestWithParam<
       _sendPtr[i] = _sendBuffer[i]->data();
       if (allocateRecvBuffer) _recvPtr[i] = _recvBuffer[i]->data();
     }
+#if UCXX_ENABLE_RMM
+    if (_bufferType == ucxx::BufferType::RMM) { rmm::cuda_stream_default.synchronize(); }
+#endif
   }
 
   void copyResults()
@@ -141,6 +144,9 @@ class RequestTest : public ::testing::TestWithParam<
 #endif
       }
     }
+#if UCXX_ENABLE_RMM
+    if (_bufferType == ucxx::BufferType::RMM) { rmm::cuda_stream_default.synchronize(); }
+#endif
   }
 };