How to back-order specific values

An int tensor his dimension is [batch,n,m], I want to move his element of last dim whose element is not -1 forward and -1 back.Is it possible to achieve the same result without using a for loop, as a for loop would be very time consuming? Here is a 2D example

import torch
tensor = torch.tensor([[1, -1, 2, -1],
                       [7, -1, 5, 6],
                       [0, 7, -1, 8]])

for i in range(tensor.shape[0]):
    tensor[i] = torch.cat((tensor[i][tensor[i] != -1], tensor[i][tensor[i] == -1]))
print(tensor)
#tensor([[ 1,  2, -1, -1],
#       [ 7,  5,  6, -1],
#      [ 0,  7,  8, -1]])

Anytime you want to do a batch logic argument on a tensor, you can create a boolean mask. For example:

“whose element is not -1” could be determined elementwise with:

mask = tensor!=-1

And then we can reorder those with:

x = torch.stack([torch.cat([tensor[_,:][mask[_,:]], tensor[_,:][~mask[_,:]]], dim=0) for _ in range(tensor.size(0))])

In this case, the for loop gets processed in parallel for each row.