Update categorical.py

[cyprienc]

Setting dtype in jax.nn.one_hot calls to avoid a return dtype different from the parameters' dtypes.

[franrruiz]

Thank you @cyprienc for your contribution to Distrax!