Note that torch.multinomial might be better if you want to stay in PyTorch.
I would recommend to avoid mixing random number generators from different frameworks if you can.
(To me using two RNGs from two different frameworks falls under “being clever” and my rule of thumb here is “don’t try to be clever with random number generation unless you absolutely know what you are doing 2”… The linked anecdote is not immediately applicable, but it is a trap I have seen people fall into. The other day the big surprise was that PyTorch dataloaders didn’t initialize the numpy seed in workers, that has been fixed in recent PyTorch, but I’d venture that it is a good habit to stay with PyTorch unless you have to use something else.)
Best regards
Thomas