torch.accelerator.set_stream

torch.accelerator.set_stream(stream)

将当前流设置为给定流。

参数

stream (torch.Stream) – 给定流，必须与当前的 accelerator 设备类型匹配。

注意

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