I dont think there’s a pytorch equivalent, but you could look at using apply_() 1.3k and then calling   nonzero() to gather the indices.
Also, apply_() is an inplace operation and only works on CPU tensors.