c10_npu::NPUStreamGuard::NPUStreamGuard

产品支持情况

产品	是否支持
Atlas A3 训练系列产品	√
Atlas A2 训练系列产品	√

功能说明

构造函数，创建一个流guard。将当前的NPU设备设置为与传入流关联的设备，并将该设备上的当前NPU流设置为传入流。NPUStreamGuard在析构时会自动恢复构造前的NPU设备和流，以确保作用域结束后恢复原始上下文。

定义文件

torch_npu\csrc\core\npu\NPUGuard.h

函数原型

c10_npu::NPUStreamGuard::NPUStreamGuard(c10::Stream stream)

参数说明

stream (c10::Stream)：必选参数，表示guard保障的流。

返回值说明

无

约束说明

无