Solved by ptrblck in post #2 
              
              
                I’m not sure to completely understand the use case, but this code snippet should work, if you would like to sample either the element from a or b based on the probabilities: 
L = 10
a = torch.arange(L)
b = torch.arange(L, 2*L)

p_a = 0.2
probs = torch.empty(1, L).uniform_()
idx = (probs >= p_a).to(t…