torch.accelerator.set_stream

torch.accelerator.set_stream(stream)

设置当前流到指定的流。

参数:

stream (torch.Stream) – 指定的流，必须与当前加速器设备类型匹配。

注意

此函数将当前设备索引设置为给定流的设备索引。