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

1

I want to copy a 2-d torch tensor to a destination tensor containing only values until the first occurrence of 202 value and zero for the rest of the items like this:

source_t=tensor[[101,2001,2034,1045,202,3454,3453,1234,202]
         ,[101,1999,2808,202,17658,3454,202,0,0]
         ,[101,2012,3832,4027,3454,202,3454,9987,202]]

destination_t=tensor[[101,2001,2034,1045,202,0,0,0,0]
                    ,[101,1999,2808,202,0,0,0,0,0]
                    ,[101,2012,3832,4027,3454,202,0,0,0]]

how can i do it?

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)