stream_guard¶

paddle.device.cuda. stream_guard ( stream ) [source]

Notes:: This API only supports dygraph mode currently.

A context manager that specifies the current stream context by the given stream.

Parameters: stream (paddle.device.cuda.Stream) – the selected stream. If stream is None, just yield.

Examples

           # 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

          