How to get 2d sub tensor with specific value condition from an original tensor in pytorch?

2

Hi,

I am not sure if it is the most efficient, but the following works:

mask = source_t.eq(202).flip(1).cumsum(dim=1).flip(1)
mask = mask.eq(mask.select(1, 0).unsqueeze(1))
print(mask) # mask of the source values to keep

res = mask * source_t
print(res)