1

My question is nearly identical to this one, with the notable difference of being in PyTorch. I would prefer not to use the Numpy solution as this would involve moving data back to the CPU. I see that, as with Numpy, PyTorch has a nonzero function, however its where function (the solution in the Numpy thread I linked) has behavior different from Numpy's. The behavior I want is an is_zero() function as follows:

>>> arr.nonzero()
tensor([[0, 1],
        [1, 0]])  
>>> arr.is_zero()
tensor([[0, 0],
        [1, 1]])
devck
  • 331
  • 2
  • 8

1 Answers1

2

You can make a boolean mask and then call nonzero():

(arr == 0).nonzero()

For instance:

arr = torch.randint(high=2, size=(3, 3))

tensor([[1, 1, 0],  # (0, 2)
        [1, 1, 0],  # (1, 2)
        [1, 0, 0]]) # (2, 1) and (2, 2)

(arr == 0).nonzero()

tensor([[0, 2],
        [1, 2],
        [2, 1],
        [2, 2]])
Nicolas Gervais
  • 33,817
  • 13
  • 115
  • 143