Replace all nonzero values by zero and all zero values by a specific value

Viewed 29294

I have a 3d tensor which contains some zero and nonzero values. I want to replace all nonzero values by zero and zero values by a specific value. How can I do that?

3 Answers

Use

torch.where(<your_tensor> != 0, <tensor with zeroz>, <tensor with the value>)

Example:

>>> x = torch.randn(3, 2)
>>> y = torch.ones(3, 2)
>>> x
tensor([[-0.4620,  0.3139],
         [ 0.3898, -0.7197],
         [ 0.0478, -0.1657]])
>>> torch.where(x > 0, x, y)
Tensor([[ 1.0000,  0.3139],
        [ 0.3898,  1.0000],
        [ 0.0478,  1.0000]])

See more at: https://pytorch.org/docs/stable/generated/torch.where.html

This can be done without cloning the tensor and using indices of zero and non-zero values:

zero_indices = tensor == 0
non_zero_indices = tensor != 0
tensor[non_zero_indices] = 0
tensor[zero_indices] = value
Related