torch.nonzero
torch.nonzero(input, *, out=None, as_tuple=False) -> LongTensor or tuple of LongTensors
Note:
 "torch.nonzero(..., as_tuple=False)" (default) returns a 2-D
 tensor where each row is the index for a nonzero
 value."torch.nonzero(..., as_tuple=True)" returns a tuple of 1-D
 index tensors, allowing for advanced indexing, so
 "x[x.nonzero(as_tuple=True)]" gives all nonzero values of tensor
 "x". Of the returned tuple, each index tensor contains nonzero
 indices for a certain dimension.See below for more details on the
 two behaviors.When "input" is on CUDA, "torch.nonzero()" causes
 host-device synchronization.

When "as_tuple" is "False" (default):
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).