torch.cuda.stream

torch.cuda.stream(stream) [исходный код]

Обёртка над контекстным менеджером StreamContext, выбирающим заданный поток.

Параметры:

stream (Stream) – выбранный поток. Этот менеджер ничего не делает, если это None.

Тип возвращаемого значения:

StreamContext

Примечание

В eager-режиме stream имеет тип класса Stream, а в JIT является объектом пользовательского класса torch.classes.cuda.Stream.