Hi,
I don’t think there is I’m afraid.
The best way I can think of to do this is:
out = torch.zeros_like(A)
n = A.size(1)
out.scatter_(1, n - 1 - B, 1) # Work in reverse order
out = out.cumsum(dim=1)
out.flip(1)