tensordict.nn.distributions.OneHotCategorical

class tensordict.nn.distributions.OneHotCategorical(logits: Optional[Tensor] = None, probs: Optional[Tensor] = 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.