torch.cuda.get_stream_from_external

torch.cuda.get_stream_from_external(data_ptr, device=None)

返回從外部分配的 CUDA 流建立的 Stream 物件。

此函式用於封裝在其他庫中分配的流，以便於資料交換和多庫互動。

注意

此函式不管理流的生命週期，使用者有責任在使用此返回的流時保持引用的流處於活動狀態。

引數

• data_ptr (int) – 外部分配的 cudaStream_t 值的整數表示。

• device (torch.device 或 int, 可選) – 流最初分配在哪個裝置上。如果裝置指定不正確，後續使用此流的啟動可能會失敗。

返回型別

Stream