```
import torch
# Assuming that I have a tensor a.
a = torch.Tensor([-1,4,6,3,-1,-1])
# Then I need to select a random element(not -1) from a.
# How to do that efficiently?
```

This code should work:

```
a = torch.Tensor([-1,4,6,3,-1,-1])
valid_idx = (a!=-1).nonzero().view(-1)
choice = torch.multinomial(valid_idx.float(), 1)
a[valid_idx[choice]]
```

I’m not sure if it’s more efficient than just calling `torch.randint(0, a.size(0), (1,))`

in a while loop to skip the `-1`

entries. It might depend on the ratio between valid and invalid entries.

1 Like

Hi, ptrblck

```
import torch
# Assuming that I have a tensor a.
a = torch.Tensor([[3,4, -1, -1,5,6],
[-1,5,-1,1,-1,-1],
[-1,4,4,2,3,0]])
# Then I need to select a random element(not -1) from a for each line
# How to do that efficiently?
# Eg.
# a_sample = [[4],[1],[3]]
```

In that case, this should work:

```
a = torch.Tensor([[3,4, -1, -1,5,6],
[-1,5,-1,1,-1,-1],
[-1,4,4,2,3,0]])
valid_idx = (a!=-1).nonzero()
choice = torch.multinomial(torch.arange(valid_idx.size(0)).float(), 1)
a[valid_idx[choice].squeeze().chunk(2)]
```

EDIT: Sorry, I didn’t realized you would like to sample from each line.

Will update the code in a minute.

EDIT2: I’m not sure, if we can avoid using a loop in this case:

```
valid_idx = (a!=-1).nonzero()
unique_rows = valid_idx[:, 0].unique()
valid_row_idx = [valid_idx[valid_idx[:, 0] == u] for u in unique_rows]
ret = []
for v in valid_row_idx:
choice = torch.multinomial(torch.arange(v.size(0)).float(), 1)
ret.append(a[v[choice].squeeze().chunk(2)])
ret = torch.stack(ret)
```

Hi, thank you for a great code snippet.

I found an error case where your code returns a Runtime Error.

When there is only one element in `valid_idx`

and that element is at the 0th position in `a`

it returns the following error:

```
RuntimeError: invalid argument 2: invalid multinomial distribution (sum of probabilities <= 0) at /Users/distiller/project/conda/conda-bld/pytorch_1556653464916/work/aten/src/TH/generic/THTensorRandom.cpp:343
```

My reproducible code is as follows:

```
a = torch.Tensor([-1,4,6,3])
indices = (a == -1).nonzero().view(-1)
choice = torch.multinomial(indices.float(), 1)
```

Your second suggestion works fine with this case.