Pytorch : Choosing randomly after softmax

Question

I use a model that ends with a softmax layer. It returns a 10000 * n tensor, and then I have to choose a number between 0 and n - 1 depending on the probability distribution for each row (so I need to return 10000 numbers in total).

I tried to use random.choices but it's not working with tensors and it will probably take a lot of time to make a list from each row.

For example, if a row is [0.1, 0.2, 0.1, 0.6], I want to return 0 with probability 0.1, 1 with probability 0.2, and so on.