Give a 2D matrix of integers, I would like to get the index of the row that has the corresponding value. For example, if I have a matrix A = [[1,2], [2,3], [3,4]] and if I have the value for column 1 =3, column 2=4, which leads to [3,4]. Then this function should return index number 2 from matrix A given that the third row corresponds to [3,4]. How can I do this in PyTorch?
You could try the following snippet:
a = torch.FloatTensor([[1, 2], [2, 3], [3, 4]]) b = torch.FloatTensor([[3, 4]]) (a == b).nonzero()
Thank you. But this returns the following:
2 0 2 1 [torch.LongTensor of size 2x2]
Do you know if it is possible to return the row index, which is 2 in our case?
How about this?
number1 = 3 # look for 3 in column 1 number2 = 4 column1 = A[:,0] column2 = A[:,1] _, index1 = torch.min(torch.abs(column1 - number1), 0) _, index2 = torch.min(torch.abs(column2 - number2), 0)
N.B. This isn’t differentiable.