stream_guard¶
可以切换当前的 stream 为输入指定的 stream。
注解
该 API 目前仅支持动态图模式。
参数¶
stream (paddle.device.Stream) - 指定的 stream。如果为 None,则不进行 stream 切换。
代码示例¶
# required: custom_device
import paddle
paddle.set_device('custom_cpu')
s = paddle.device.Stream()
data1 = paddle.ones(shape=[20])
data2 = paddle.ones(shape=[20])
data3 = data1 + data2
with paddle.device.stream_guard(s):
s.wait_stream(paddle.device.default_stream())
data4 = data1 + data3