Hi Luiz!



 Charamba:

Hi Guys,
I would like to remove zero values of a tensor
…
In an output tensor like this:
out_x = torch.tensor([[[[  2.,  20., 250., 250.,  20.,   2.,   0.,   0.,   0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20.,  20., 250., 250., 20.,  20.,   2.,    0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20., 250., 250.,  20., 20.,   2.,   0.,   0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.],
                       [  2.,  20.,  20., 250., 250.,  20.,  20.,   2., 0.,   0.,   0.,
                          0.,   0.,   0.,   0.,   0.]]]])



At the cost of  n log (n) time complexity, you can use argsort()* and
then use gather() to index back into your input tensor:
>>> import torch
>>> torch.__version__
'1.9.0'
>>> in_x = torch.tensor([[[[0.,   0.,   2.,  20., 250.,   0.,   0.,   0.,   0.,   0., 250.,
...                       20.,   2.,   0.,   0.,   0.],
...                       [0.,   0.,   2.,  20.,   0.,   0.,  20., 250.,   0.,   0., 250.,
...                       20.,  20.,   2.,   0.,   0.],
...                       [0.,   0.,   0.,   0.,   0.,   0.,   2.,  20., 250.,   0., 250.,
...                       20.,  20.,   2.,   0.,   0.],
...                       [0.,   2.,  20.,   0.,  20., 250.,   0., 250.,  20.,   0.,  20.,
...                        2.,   0.,   0.,   0.,   0.]]]])
>>>
>>> in_x.gather (3, (in_x == 0.0).sort (dim = 3, stable = True)[1])
tensor([[[[  2.,  20., 250., 250.,  20.,   2.,   0.,   0.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.],
          [  2.,  20.,  20., 250., 250.,  20.,  20.,   2.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.],
          [  2.,  20., 250., 250.,  20.,  20.,   2.,   0.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.],
          [  2.,  20.,  20., 250., 250.,  20.,  20.,   2.,   0.,   0.,   0.,
             0.,   0.,   0.,   0.,   0.]]]])

*)  Except you have to use sort() because argsort() doesn’t support
stable = True).
Best.
K. Frank