mindspore.runtime.set_cur_stream

View Source On Gitee
mindspore.runtime.set_cur_stream(stream)[source]

Sets the current stream.This is a wrapper API to set the stream. Usage of this function is discouraged in favor of the stream context manager.

Parameters

stream (Stream) – selected stream. This function is a no-op if this argument is None.

Raises

TypeError – If 'stream' is neither a mindspore.runtime.Stream nor a None.

Examples

>>> import mindspore as ms
>>> ms.set_device("Ascend", 0)
>>> cur_stream = ms.runtime.current_stream()
>>> assert cur_stream == ms.runtime.default_stream()
>>> s1 = ms.runtime.Stream()
>>> ms.runtime.set_cur_stream(s1)
>>> assert ms.runtime.current_stream() == s1
>>> ms.runtime.set_cur_stream(ms.runtime.default_stream())