reference/generated/tensordict.TensorDictBase.numpy

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.numpy

TensorDictBase.numpy() → ndarray | dict[str, Any]

Returns a nested dict of NumPy arrays with the entries of the tensordict.

Each tensor becomes a NumPy array with torch.Tensor.numpy(), so the arrays share memory with the tensors, and a tensor that requires gradients or is not on the CPU raises a RuntimeError or TypeError: call detach() or cpu() first. A nested tensor becomes a tuple of arrays. Non-tensor data is returned as its Python value. The None fields of tensorclasses are dropped.

See also

to_struct_array() to convert to a NumPy structured array, and to_dict().

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> data = TensorDict({"a": {"b": torch.zeros(()), "c": "a string!"}})
>>> data.numpy()
{'a': {'b': array(0., dtype=float32), 'c': 'a string!'}}

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