On line 169 of README, dist = ntorch.distributions.Categorical(logits=params, logit_dim="logits") is incorrect, in that the keyword is dim_logitas seen on line 27 of namedtensor/distributions/distributions.py. While a simple fix would be changing the README, I actually prefer logit_dim over dim_logit so would prefer changing the keyword. This, however, has issues with backwards compatibility.