@melike  yes, in this case i want to remove ‘2’ in x, ‘1’ is the mask ID.
x = torch.tensor([[1., 2., 2., 2., 3.],
        [1., 2., 4., 3., 2.]])

#convert to:

tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])