torch.argwhere
torch.argwhere(input) -> Tensor
Returns a tensor containing the indices of all non-zero elements of
   "input".  Each row in the result contains the indices of a non-zero
   element in "input". The result is sorted lexicographically, with
   the last index changing the fastest (C-style).
If "input" has n dimensions, then the resulting indices tensor
   "out" is of size (z \times n), where z is the total number of non-
   zero elements in the "input" tensor.
Note:
 This function is similar to NumPy's *argwhere*.When "input" is on
 CUDA, this function causes host-device synchronization.

Parameters:
      {input} --
Example:
  >>> t = torch.tensor([1, 0, 1])
  >>> torch.argwhere(t)
  tensor([[0],
          [2]])
  >>> t = torch.tensor([[1, 0, 1], [0, 1, 1]])
  >>> torch.argwhere(t)
  tensor([[0, 0],
          [0, 2],
          [1, 1],
          [1, 2]])
