torch.multinomial
torch.multinomial(input, num_samples, replacement=False, *, generator=None, out=None) -> LongTensor
Returns a tensor where each row contains "num_samples" indices
   sampled from the multinomial probability distribution located in
   the corresponding row of tensor "input".
Note:
 The rows of "input" do not need to sum to one (in which case we
 use the values as weights), but must be non-negative, finite and
 have a non-zero sum.

Indices are ordered from left to right according to when each was
   sampled (first samples are placed in first column).
If "input" is a vector, "out" is a vector of size "num_samples".
If "input" is a matrix with m rows, "out" is an matrix of shape
   (m \times \text{num_samples}).
If replacement is "True", samples are drawn with replacement.
If not, they are drawn without replacement, which means that when a
   sample index is drawn for a row, it cannot be drawn again for that
   row.
Note: