torch.cuda.stream

torch.cuda.stream(stream)

包裝 StreamContext 上下文管理器，以選擇指定的 stream。

參數

stream (Stream) – 要選擇的 stream。如果為 None，則此管理器不執行任何操作。

回傳類型

StreamContext

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