From d478cda217a9bcc5a6c89f72f68b1cd9daf9cab5 Mon Sep 17 00:00:00 2001
From: linmin <mavenlin@gmail.com>
Date: Thu, 26 Oct 2023 14:42:40 +0800
Subject: [PATCH] fix segfault issue #283

---
 envpool/core/xla.h | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/envpool/core/xla.h b/envpool/core/xla.h
index e70fbb98..f4d9afff 100644
--- a/envpool/core/xla.h
+++ b/envpool/core/xla.h
@@ -74,8 +74,8 @@ Array GpuBufferToArray(cudaStream_t stream, const void* buffer,
     spec = spec.Batch(batch_size);
   }
   Array ret(spec);
-  cudaMemcpy(ret.Data(), buffer, ret.size * ret.element_size,
-             cudaMemcpyDeviceToHost);
+  cudaMemcpyAsync(ret.Data(), buffer, ret.size * ret.element_size,
+                  cudaMemcpyDeviceToHost, stream);
   return ret;
 }
 
@@ -161,6 +161,7 @@ struct XlaSend {
            ...);
         },
         action_spec);
+    cudaStreamSynchronize(stream);
     envpool->Send(action);
   }
 };