pytorch中where方法介绍

 

一、具体用法

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

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

 

二、示例

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))