Find index where a sub-tensor does not equal to a given tensor in Pytorch

user_13470983 0
Viewed 143

I have a tensor, for example,

a = [[15,30,0,2], [-1,-1,-1,-1], [10, 20, 40, 60], [-1,-1,-1,-1]]

which has the shape (4,4). How can I find the index where a specific sub-tensor

[-1,-1,-1,-1]

that doesn't appear using PyTorch. The expected output I want to get is

[0,2]
pytorch python
2 Answers

You can compare the elements for each row of the tensor using torch.any(), and then use .nonzero() and .flatten() to generate the indices:

torch.any(a != torch.Tensor([-1, -1, -1, -1]), axis=1).nonzero().flatten()

For example,

import torch

a = torch.Tensor([[15,30,0,2], [-1,-1,-1,-1], [10, 20, 40, 60], [-1,-1,-1,-1]])
result = torch.any(a != torch.Tensor([-1, -1, -1, -1]), axis=1).nonzero().flatten()
print(result)

outputs:

tensor([0, 2])
user_177698150

You can also use where or nonzero:

a = torch.Tensor([[15,30,0,2], [-1,-1,-1,-1], [10, 20, 40, 60], [-1,-1,-1,-1]])

b = torch.Tensor([-1,-1,-1,-1])

result = torch.where(a != b)[0].unique()

result = torch.nonzero(a != b, as_tuple=True)[0].unique()

print(result)
user_167331010
Related