tensordict.nn.distributions.OneHotCategorical

class tensordict.nn.distributions.OneHotCategorical(logits: torch.Tensor | None = None, probs: torch.Tensor | None = None, **kwargs)

One-hot categorical distribution.

This class behaves excacly as torch.distributions.Categorical except that it reads and produces one-hot encodings of the discrete tensors.