torchwrench.nn.functional.activation module¶

torchwrench.nn.functional.activation.log_softmax_multidim(x: Tensor, *, dims: Iterable[int] | None = (-1,)) → Tensor[source]¶

A multi-dimensional version of torch.log_softmax along multiple dimensions at the same time.

torchwrench.nn.functional.activation.softmax_multidim(x: Tensor, *, dims: Iterable[int] | None = (-1,)) → Tensor[source]¶

A multi-dimensional version of torch.softmax along multiple dimensions at the same time.