How to select random index with condition?

input_boolean_mask= (torch.FloatTensor(10).uniform_()>0.8) # shape 10  boolean tensor
#tensor([False,  True, False, False, False, False, False, False,  True, False])

I want to random select “n” indexes only among the values corresponding to the conditions I want( for above example, True value).

How can I do this? Thanks.

Something like this should work:

input_boolean_mask = (torch.empty(10).uniform_()>0.8)
print(input_boolean_mask)
# tensor([ True, False, False, False, False,  True,  True, False, False,  True])

# get indices of True values
idx = input_boolean_mask.nonzero()
print(idx)
# tensor([[0],
#         [5],
#         [6],
#         [9]])

# select n indices
n = 2
ret = idx[:n]
print(ret)
# tensor([[0],
#         [5]])

I wanted to select random n indices among True values. For example, select random n selection among indices (0,5,6,9).
I solved problem with torch.Tensor.multinomial (torch.Tensor.multinomial — PyTorch 1.13 documentation)

input_boolean_mask = (torch.empty(10).uniform_()>0.8).long()
input_boolean_mask=input_boolean_mask.float()
print(input_boolean_mask)
#tensor([0, 1, 0, 0, 0, 0, 1, 0, 1, 0])
indices=input_boolean_mask.multinomial(2,replacement=False)
print (indices)
#tensor([1,8])

If this way is inefficient and you have suggestion, please let me know.

Thanks.