Pytorch, get the index of the first 0 in as mask?

Question

I have a tensor that looks like: (1, 1, 1, 1, 1, 1, 1, 1, 0, 0). I want to get the index where the first zero appears. What would the be best way do this?

You can adapt this answer to your case https://stackoverflow.com/a/60817512/8440629 — Tajinder Singh, Jul 20 '22 at 10:19
if it's guaranteed to be contiguous 1's followed by contiguous 0's then the sum will give you the index of the first 0. — Archie Gertsman, Nov 06 '22 at 18:31

score 1 · Accepted Answer · answered Jul 20 '22 at 10:16

1

Not the best usage of argmin but it should work here I think:

>>> torch.tensor([1, 1, 1, 1, 1, 1, 1, 1, 0, 0]).argmin()
tensor(8)

answered Jul 20 '22 at 10:16

Serhii Maksymenko

259
1
7

Oh, I'm very silly, the mask is actually a boolean mask. I'll see if `argmin` still works. – Foobar Jul 20 '22 at 10:29
it works only if the element is minimum. – bpfrd Jul 20 '22 at 10:45
This: `mask = mask.to(torch.int16)`, solved the issue – Foobar Jul 20 '22 at 10:52
solution from @bpfrd is more universal because it will work with any target value (not only zero) – Serhii Maksymenko Jul 20 '22 at 12:12

score 1 · Answer 2 · answered Jul 20 '22 at 10:44

1

try this:

your_target_value = 0
your_tensor = torch.tensor([1, 1, 1, 1, 1, 1, 1, 1, 0, 0])

(your_tensor == your_target_value).nonzero()[0] #first element

Output:

tensor([8])

answered Jul 20 '22 at 10:44

bpfrd

945
3
11

Pytorch, get the index of the first 0 in as mask?

2 Answers2