torch.accelerator.set_stream

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

Устанавливает текущий поток в указанный поток.

Параметры:

stream (torch.Stream) – указанный поток, тип устройства которого должен соответствовать типу текущего устройства ускорителя.

Примечание

Эта функция установит индекс текущего устройства равным индексу устройства указанного потока.