torch.mtia.stream

torch.mtia.stream(stream)

將 Context-manager StreamContext 包裝起來，以選擇指定的 stream。

參數

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

回傳類型

StreamContext

..Note:: 在 Eager 模式下，stream 的類型為 Stream 類別，而在 JIT 中，則不支援 torch.mtia.stream