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)