vec_len = 5

a = torch.rand(vec_len, 1)
b = torch.rand(vec_len, 1)
c = torch.rand(vec_len, 1)

aa = torch.tile(a, (1, 3))
bb = torch.tile(b, (1, 3))
cc = torch.tile(c, (1, 3))

index_tensor = torch.Tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]])
index_tensor_ = torch.tile(index_tensor.view(-1, 3, 3), (vec_len, 1, 1))

stacked = torch.stack((aa, bb, cc), dim=2).type(torch.double)
result = torch.where(index_tensor_ > 0, stacked, 0.)

This is the best among what I know
Any idea?