- torch.nn.functional.softmax(input, dim=None, _stacklevel=3, dtype=None)¶
Applies a softmax function.
Softmax is defined as:
It is applied to all slices along dim, and will re-scale them so that the elements lie in the range [0, 1] and sum to 1.
Softmaxfor more details.
input (Tensor) – input
dim (int) – A dimension along which softmax will be computed.
torch.dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is casted to
dtypebefore the operation is performed. This is useful for preventing data type overflows. Default: None.
- Return type
This function doesn’t work directly with NLLLoss, which expects the Log to be computed between the Softmax and itself. Use log_softmax instead (it’s faster and has better numerical properties).