一、具体用法
torch.where(condition, x, y)
→ 其中 x,y 必须是Tensor类型
※ 共有三个输入参数,第一个是condition判断条件,第二个是符合条件时取设置值x,第三个是不满足条件的设置值y;
即类似于
a
n
s
=
{
x
,
c
o
n
d
i
t
i
o
n
y
,
!
c
o
n
d
i
t
i
o
n
ans=\left\{\begin{matrix} &x, &condition \\ &y, &!\ condition \end{matrix}\right.
ans={x,y,condition! condition
二、示例
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))