Solved by KFrank in post #2 
              
              
                Hi Yoda! 

As I understand your question, you are asking for the categorical 
distribution (which is a special case of the multinomial distribution). 
torch.distributions.Categorical take batches of probs (as does 
torch.distributions.Multinomial), so you can just pass your probs 
in as-is. 
Here is…