CODEWITHSUNDEEP
pytorchtensortorch
Foobar
Foobar

Reputation: 8467

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

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?

Upvotes: 1

Views: 1760

Answers (2)

bpfrd
bpfrd

Reputation: 1025

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])

Upvotes: 3

Serhii Maksymenko
Serhii Maksymenko

Reputation: 319

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)

Upvotes: 1

Related Questions