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(torch.long)
c = torch.stack((a, b))
result = torch.gather(c, 0, idx)

print(c)
> tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
          [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])
print(probs)
> tensor([[0.8268, 0.0361, 0.2744, 0.4381, 0.7244, 0.8509, 0.5047, 0.0549, 0.2767,           0.3203]])
print(idx)
> tensor([[1, 0, 1, 1, 1, 1, 1, 0, 1, 1]])
print(result)
> tensor([[10,  1, 12, 13, 14, 15, 16,  7, 18, 19]])