I don’t find that really straight forward.
First of all you would need to calculate which is the nearest neightbor.
You can get non-zero elements by doing tensor[tensor>0].
This is already bit expensive.
But this way you are losing the position in the grid.
Getting that is hard.
In fact you could try to use sparse tensors.
https://pytorch.org/docs/stable/sparse.html
as they save the position of the non-zero elements.
However as u can see here:
  
    
    
    Dense to Sparse Tensors 
  
  
    Is there a way to convert dense tensors to sparse tensors without having to extract every value and coordinate and recreating the dense tensor as a sparse one?
  


The developer point out that there is no easy way to obtain a sparse tensor from a dense one. Which is basically the problem you are trying to solve.
If you can construct the sparse tensor you can achieve good performance.
Otherwise I would recomment to find another way to fill the non-zero elements or to relax your criteria bout what’s the nearest.
For example something very silly that would be fast is to run a convolution (as the operation takes the max of the kernel). If you matrix is really sparse and ur kernel is big enough you would me taking some sort of nearest neighbor.
hope it helps