TensorDictBase.softmax
-
TensorDictBase.softmax(dim: int, dtype: dtype | None = None)
Apply a softmax function to the tensordict elements.
- Parameters:
dim (int or tuple of ints) – A tensordict dimension along which softmax will be computed.
dtype (dtype, optional) – the desired data type of returned tensor.
If specified, the input tensor is cast to dtype before the operation is performed.
This is useful for preventing data type overflows.
Every tensordict class and every tensorclass has this method. For the
other methods, see Methods by task.
