mindspore.hal.StreamCtx

查看源文件
class mindspore.hal.StreamCtx(ctx_stream)[源代码]

上下文管理器,用于选择给定的流。

上下文范围内的所有算子,都将在指定流上执行。

参数:
  • ctx_stream (Stream) - 指定的流。如果是 None ,这个上下文管理器无操作。

异常:

样例:

>>> import mindspore as ms
>>> import numpy as np
>>> from mindspore import Tensor, ops
>>> a = Tensor(np.ones([1024, 2048]), ms.float32)
>>> b = Tensor(np.ones([2048, 4096]), ms.float32)
>>> s1 = ms.hal.Stream()
>>> with ms.hal.StreamCtx(s1):
...     c = ops.matmul(a, b)
>>> ms.hal.synchronize()
>>> assert s1.query()