torch.mtia.stream

torch.mtia.stream(stream)

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

引數

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

返回型別

StreamContext

注意

在 eager 模式下，stream 的型別為 Stream 類，而在 JIT 中，它不支援 torch.mtia.stream。