stream_guard¶
可以切换当前的 CUDA stream 为输入指定的 stream。
参数¶
stream (paddle.device.cuda.Stream) - 指定的 CUDA stream。如果为 None,则不进行 stream 流切换。
代码示例¶
# required: gpu
import paddle
s = paddle.device.cuda.Stream()
data1 = paddle.ones(shape=[20])
data2 = paddle.ones(shape=[20])
with paddle.device.cuda.stream_guard(s):
data3 = data1 + data2