Sampling from a tensor in Torch

I wanted to know whether we could sample elements of a tensor given a probability distribution of the tensor.
In numpy we would do something like this:

a = np.array([1,2,3,4])
b = np.random.choice(a, p=np.array([0.1, 0.1, 0.1, 0.7]))

In torch I would like to have the array a and p to be of torch.tensor. The numpy function works well with CPU torch tensors translated to numpy arrays, but unfortunately, using GPUs the trick fails.

Might be helpful? Torch equivalent of numpy.random.choice?

Have seen, did not do the trick.

Hi !
If you want to do :

a_numpy = np.array([1, 2, 3, 4])
p_numpy = np.array([0.1, 0.1, 0.1, 0.7])
b_numpy = np.random.choice(a_numpy, p=p_numpy)

You can use torch.multinomial :

a = torch.tensor([1, 2, 3, 4])
p = torch.tensor([0.1, 0.1, 0.1, 0.7])
index = p.multinomial(num_samples=1, replacement=True)
b = a[index]

Here I used replacement=True because it is the default behavior of numpy.choice, even though is makes no difference when num_samples=1