torch.cuda.stream

torch.cuda.stream(stream)

围绕选择给定流的上下文管理器 StreamContext 进行包装。

Parameters

流 (流) – 选择的流。如果为None，则此管理器不执行任何操作。

Return type

流上下文

..注意:: 在急切模式下，流是Stream类的类型，而在JIT中，它是自定义类torch.classes.cuda.Stream的对象。