torch.accelerator.set_stream

torch.accelerator.set_stream(stream)

Set the current stream to a given stream.

Parameters

stream (torch.Stream) – a given stream that must match the current accelerator device type.

Note

This function will set the current device index to the device index of the given stream.