Select rows of the tensor whose first element is equal to some value

#1

Given a tensor like

a =  1  1  1
     1  1  1
     1  1  1
     2  1  1
     2  1  1

I would like to obtain all rows starting with 2, i.e. I would like to obtain

b = 2 1 1 
    2 1 1

Indexing using a[a[:, 0]==2] doesn’t seem to work and index_select requires indices to be known and so does gather.

1 Like
#2

a[(a[:, 0] == 2).nonzero().squeeze(1)]

2 Likes
#3

Thanks! What if I need two criteria columns at once? for example, select rows where column A equals 2 and column B equals 5 ?

#4

Hi Hendy, did you find a solution for this?

#5

Hi @jekbradbury Could you please also kind respond to this question? Thanks a lot!