Updating pytorch tensor by matching multiple boolean conditions on the original values

user_19205811 0
Viewed 60

I currently am updating values in a Pytorch tensor using multiple OR conditions:

>>> import torch

>>> my_tensor = torch.tensor([0, 1, 2, 3, 4, 5])

>>> condition = ((my_tensor==1) | (my_tensor==4) | (my_tensor==5))
>>> my_tensor[condition] = 0

>>> my_tensor
[0, 0, 2, 3, 0, 0]

My list of conditions is much longer than the toy example above. Can the condition operator match a list? If not, what is the best solution?

tensor pytorch python
1 Answers

How about using a combination of torch.where and torch.isin like below:

>>> torch.where(torch.isin(my_tensor, torch.tensor([1,4,5])), 0, my_tensor)
tensor([0, 0, 2, 3, 0, 0])

Update, Second approach: We can use torch.reshape.

(torch.isin not exist in pytorch==1.9.1 as you say in comment)

>>> mask = (my_tensor == torch.reshape(torch.tensor([1,4,5]), (-1,1))).any(0)
>>> torch.where(mask, 0, my_tensor)
user_17405770
Related