I have a 2-D Tensor A of shape = [ BS, n*c ]

I need a 2-D Tensor A_out of shape = [ BS, (n-1)*c ]

Essentially, I want to delete c values from each row. These c values will come from different columns for each row.

These specific columns will be decided by the below tensor L:

Shape- [ BS, 1 ] and Values can be integers from ( 0, n-1 )

I want to remove the values between indices c * L[i] and c * ( L[i] + 1 ) for row i

The best I could come up with is as follows. If anyone can help with removing the for loop in this, it’ll be really appreciated:

A_out = torch.zeros( BS, (n-1) * c)

for i in range( BS ):

A_out[i] = torch.cat( [ A[ i , :c * L[i] ], A[ i , c * ( L[i]+1 ):] ] )

(‘gather’ seemed like the appropriate option at first, but creating the 2-D tensor of indexes to gather using L also seemed equally costly i.e. for loop needed)