I have a tensor A of shape m*n and a long tensor B of shape m*1, each element in B is in [0, n-1] indicating the last unmasked element in each row of A, and then I can return a mask of m*n for A based on B.
For example, the shape of A is 4*3, and B is [[2], [1], [0], [1]], then I want to return a mask as follows,
[[1, 1, 1],
[1, 1, 0],
[1, 0, 0],
[1, 1, 0]]

Is there any similar API provided by PyTorch?