Applies a softmax function.
(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: NULL.