Respect the current stream on D2D copy
@@ -515,7 +515,9 @@ cdef class ndarray:
# it is recommended to set the current device to the device
# where the src data is physically located.
runtime.setDevice(self.data.device_id)
- newarray.data.copy_from_device(x.data, x.nbytes)
+ try:
+ newarray.data.copy_from_device_async(x.data, x.nbytes)
+ finally:
if runtime._is_hip_environment:
runtime.setDevice(dev_id)
return newarray
