pytorch中where方法介绍

一、具体用法

       torch.where(condition, x, y)→ 其中 x,y 必须是Tensor类型

※ 共有三个输入参数，第一个是condition判断条件，第二个是符合条件时取设置值x，第三个是不满足条件的设置值y；
                     即类似于 $ans=\left\{\begin{matrix} &x, &condition \\ &y, &!\ condition \end{matrix}\right.$

二、示例

import torch

ans = torch.randn(3, 2)
x = torch.randn(3, 2)
y = torch.randn(3, 2)
print(ans)
print(x)
print(y)
print(torch.where(ans > 0, x, y))

Out：
              

三、设置值必须是tensor类型

import torch

ans = torch.randn(3, 2)

print(ans)
print(torch.where(ans > 0, 0, 1))