stream_guard

paddle.device.cuda. stream_guard ( stream ) [源代码]

可以切换当前的 CUDA stream 为输入指定的 stream。

注解

该 API 目前仅支持动态图模式。

参数

  • stream (paddle.device.cuda.Stream) - 指定的 CUDA stream。如果为 None,则不进行 stream 流切换。

代码示例

>>> import paddle
>>> paddle.device.set_device('gpu')

>>> 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