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
deviceis not a cuda device.
Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.