And you should check shape[-1], just in case you have a batched Categorical RV
1 Like
And you should check shape[-1], just in case you have a batched Categorical RV