reference/generated/tensordict.TensorDictBase.cuda

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.cuda

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

Returns the tensordict on a CUDA device, as to() does.

A tensordict that is already on that device is returned as it is.

Parameters:

• device (int, str or device, optional) – the CUDA device. An integer (or a string of digits) is a device index. Defaults to None, which uses the current CUDA device.

• **kwargs – passed to to(), for example non_blocking.

Raises:

RuntimeError – if device is not a CUDA device.

Examples

The example needs a CUDA device, so it is not run when the documentation is tested.

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), batch_size=[3])
>>> td.cuda().device  
device(type='cuda', index=0)
>>> td.cuda(1).device  
device(type='cuda', index=1)

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