Equivalent torch tensor function for “find”?

Question

In a nutshell, I would like to know if there is a tensor command in the torch that gives me the indices of the elements in the tensor that satisfy certain criteria.

Here is the Matlab code that illustrates what I would like to do in a torch:

my_mat = magic(3); % returns a 3 by 3 matrix with the numbers 1 through 9
greater_than_fives = find(my_mat > 5); % find indices of all values greater than 5, the " > 5" is a logical elementwise operator that returns a matrix of all 0 and 1 and finally the "find" command picks out the indices with a "1" in them
my_mat(greater_than_fives) = 0;  % set all values greater than 5 equal to 0 

I understand that I could do this in a torch using a for loop, but is there any equivalent to the find command for matlab that would allow me to do this more compactly?

Answer 1

x [x: gt (5)] = 0

In general, there is x: gt: lt: ge: le: eq

There is also a general: apply the tha function in an anonymous function and apply it to each element.