torch.cuda

This package adds support for CUDA tensor types, that implement the same function as CPU tensors, but they utilize GPUs for computation.

It is lazily initialized, so you can always import it, and use is_available() to determine if your system supports CUDA.

CUDA semantics has more details about working with CUDA.

torch.cuda.current_blas_handle()

Returns cublasHandle_t pointer to current cuBLAS handle

torch.cuda.current_device()

Returns the index of a currently selected device.

torch.cuda.current_stream()

Returns a currently selected Stream.

class torch.cuda.device(idx)

Context-manager that changes the selected device.

Parameters:idx (int) – device index to select. It’s a no-op if this argument is negative.

torch.cuda.device_count()

Returns the number of GPUs available.

class torch.cuda.device_of(obj)

Context-manager that changes the current device to that of given object.

You can use both tensors and storages as arguments. If a given object is not allocated on a GPU, this is a no-op.

Parameters:obj (Tensor or Storage) – object allocated on the selected device.

torch.cuda.is_available()

Returns a bool indicating if CUDA is currently available.

torch.cuda.set_device(device)

Sets the current device.

Usage of this function is discouraged in favor of device. In most cases it’s better to use CUDA_VISIBLE_DEVICES environmental variable.

Parameters:device (int) – selected device. This function is a no-op if this argument is negative.

torch.cuda.stream(stream)

Context-manager that selects a given stream.

All CUDA kernels queued within its context will be enqueued on a selected stream.

Parameters:stream (Stream) – selected stream. This manager is a no-op if it’s None.

torch.cuda.synchronize()

Waits for all kernels in all streams on current device to complete.

Communication collectives

torch.cuda.comm.broadcast(tensor, devices)

Broadcasts a tensor to a number of GPUs.

Parameters:

• tensor (Tensor) – tensor to broadcast.
• devices (Iterable) – an iterable of devices among which to broadcast. Note that it should be like (src, dst1, dst2, …), the first element of which is the source device to broadcast from.

Returns:

A tuple containing copies of the tensor, placed on devices corresponding to indices from devices.

torch.cuda.comm.reduce_add(inputs, destination=None)

Sums tensors from multiple GPUs.

All inputs should have matching shapes.

Parameters:

• inputs (Iterable[Tensor]) – an iterable of tensors to add.
• destination (int, optional) – a device on which the output will be placed (default: current device).

Returns:

A tensor containing an elementwise sum of all inputs, placed on the destination device.

torch.cuda.comm.scatter(tensor, devices, chunk_sizes=None, dim=0, streams=None)

Scatters tensor across multiple GPUs.

Parameters:

• tensor (Tensor) – tensor to scatter.
• devices (Iterable[int]) – iterable of ints, specifying among which devices the tensor should be scattered.
• chunk_sizes (Iterable[int], optional) – sizes of chunks to be placed on each device. It should match devices in length and sum to tensor.size(dim). If not specified, the tensor will be divided into equal chunks.
• dim (int, optional) – A dimension along which to chunk the tensor.

Returns:

A tuple containing chunks of the tensor, spread accross given devices.

torch.cuda.comm.gather(tensors, dim=0, destination=None)

Gathers tensors from multiple GPUs.

Tensor sizes in all dimension different than dim have to match.

Parameters:

• tensors (Iterable[Tensor]) – iterable of tensors to gather.
• dim (int) – a dimension along which the tensors will be concatenated.
• destination (int, optional) – output device (-1 means CPU, default: current device)

Returns:

A tensor located on destination device, that is a result of concatenating tensors along dim.

Streams and events

class torch.cuda.Stream

Wrapper around a CUDA stream.

Parameters:

• device (int, optional) – a device on which to allocate the Stream.
• priority (int, optional) – priority of the stream. Lower numbers represent higher priorities.

query()

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)

Records an event.

Parameters:event (Event, optional) – event to record. If not given, a new one will be allocated. Returns:Recorded event.

synchronize()

Wait for all the kernels in this stream to complete.

wait_event(event)

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

Parameters:event (Event) – an event to wait for.

wait_stream(stream)

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.

class torch.cuda.Event(enable_timing=False, blocking=False, interprocess=False, _handle=None)

Wrapper around CUDA event.

Parameters:

• enable_timing (bool) – indicates if the event should measure time (default: False)
• blocking (bool) – if true, wait() will be blocking (default: False)
• interprocess (bool) – if true, the event can be shared between processes (default: False)

elapsed_time(end_event)

Returns the time elapsed before the event was recorded.

ipc_handle()

Returns an IPC handle of this event.

query()

Checks if the event has been recorded.

Returns:A boolean indicating if the event has been recorded.

record(stream=None)

Records the event in a given stream.

synchronize()

Synchronizes with the event.

wait(stream=None)

Makes a given stream wait for the event.