I want to get the masked elements of each row in 2D tensor.  But the current method is slower :
x = torch.tensor([[1., 2., 2., 2., 3.],
        [1., 2., 4., 3., 2.]])

ones = torch.tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

masks = torch.tensor([[ True, False, False, False,  True],
        [ True, False,  True,  True, False]])

for i in range(x.size(0)):
    mask = masks[i]
    s = torch.masked_select(x[i], mask)
    ones[i][:s.size(0)] = s

print(f'ones:{ones}')

#output: 
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])



How to let the for  loop to batch?
@melike @Zhaoyi-Yan @ptrblck   Is there any solution ？Thank you!