Torch张量中与Matlab的“find”功能相当的函数是什么？

简而言之，我想知道在 torch 中是否有一个张量命令，可以给我一个在张量中满足某种条件的元素的索引。

以下是展示我想在 torch 中做的事情的 matlab 代码：

my_mat = magic(3); % 返回数字1到9的 3x3 矩阵
greater_than_fives = find(my_mat > 5); % 找到所有大于5的值的索引，"> 5" 是一个逻辑按元素运算符，返回所有0和1的矩阵，最后"find" 命令选择其中有"1"的索引
my_mat(greater_than_fives) = 0;  % 把所有大于5的值都设为0

我知道我可以用一个for 循环在 torch 中完成这个操作，但是是否有一些类似于 matlab 的 find 命令的等效命令，让我可以更简洁地完成这个操作？