Rate this Page
★ ★ ★ ★ ★

TensorDictBase.cuda#

TensorDictBase.cuda(device: int | str | device | None = None, **kwargs) → Self#

Casts a tensordict to a cuda device (if not already on it).

Parameters:

device (int, str or device, optional) – if provided, the cuda device on which the tensor should be cast. An integer is a cuda device index. Defaults to the current cuda device.

This function also supports all the keyword arguments of to().

Raises:

RuntimeError – if device is not a cuda device.

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

Tutorials that use cuda#

Devices and memory-mapped tensors

Devices and memory-mapped tensors