fix segfault issue #283

sail-sg · Oct 26, 2023 · d478cda · d478cda
1 parent 47ad258
commit d478cda
Showing 1 changed file with 3 additions and 2 deletions.
diff --git 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);
   }
 };