Shortcuts

Stream

class torch.cuda.Stream(device=None, priority=0, **kwargs)[source]

Wrapper around a CUDA stream.

A CUDA stream is a linear sequence of execution that belongs to a specific device, independent from other streams. See CUDA semantics for details.

Parameters
  • device (torch.device or int, optional) – a device on which to allocate the stream. If device is None (default) or a negative integer, this will use the current device.

  • priority (int, optional) – priority of the stream. Can be either -1 (high priority) or 0 (low priority). By default, streams have priority 0.

Note

Although CUDA versions >= 11 support more than two levels of priorities, in PyTorch, we only support two levels of priorities.

query()[source]

Checks if all the work submitted has been completed.

Returns

A boolean indicating if all kernels in this stream are completed.

record_event(event=None)[source]

Records an event.

Parameters

event (torch.cuda.Event, optional) – event to record. If not given, a new one will be allocated.

Returns

Recorded event.

synchronize()[source]

Wait for all the kernels in this stream to complete.

Note

This is a wrapper around cudaStreamSynchronize(): see CUDA Stream documentation for more info.

wait_event(event)[source]

Makes all future work submitted to the stream wait for an event.

Parameters

event (torch.cuda.Event) – an event to wait for.

Note

This is a wrapper around cudaStreamWaitEvent(): see CUDA Stream documentation for more info.

This function returns without waiting for event: only future operations are affected.

wait_stream(stream)[source]

Synchronizes with another stream.

All future work submitted to this stream will wait until all kernels submitted to a given stream at the time of call complete.

Parameters

stream (Stream) – a stream to synchronize.

Note

This function returns without waiting for currently enqueued kernels in stream: only future operations are affected.

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources