torch.where一个参数的情况

昨天查看论文源码的时候发现有使用torch.where函数，在网站上查找发现多为3参数解释，便去官网看一个参数的解释，在此记录。

以下为官方document中的解释：

它相当于as_tuple = True 时的nonzero方法。

不过，他这里只说了不为零的判断情况，其他的condition也是行得通的。

 举个简单的例子帮助理解：

假设存在一个二维数组M：

然后试图让x，y分别返回condition为tensor元素大于2的情况，从第2列开始，则代码为：

 

从结果可以看出，它是相当于对后三列的判断是否大于2，返回的x y对应位置组成的2元组即是大于2的位置的坐标信息，如第一个x y为0，1，代表0，1位置（数值为3）大于2，以此类推。