torch.cuda.stream

torch.cuda.stream(stream)

包裝選擇給定流的上下文管理器 StreamContext。

引數

stream (Stream) – 選定的流。如果為 None，則此管理器無效。

返回型別

StreamContext

注意

在 eager 模式下，stream 的型別是 Stream 類，而在 JIT 中，它是一個自定義類 torch.classes.cuda.Stream 的物件。