Suppose I have a 2D tensor. I want to get the index of the first negative element of the tensor if there are any negative elements, else get the index of the smallest element. Basically, I want to implement the following for loop over each row of a 2D tensor efficiently in pytorch. Is there a way to do this with pytorch operations in a single iteration over each row?
x = 10.0 # all numbers are guaranteed less than 10 for i in range(len(row)): if row[i] < 0: x = i break if row[i] < x: x = i