Get first negative element of row, else minimum

1

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
2

Figured it out:

            clamped = torch.clamp(rays, 5e-11, 10.0)
            minIdx = torch.argmin(clamped, dim=1)

-5e-11 is an upper bound on the negative numbers I’m interested in.