- 函數作用
torch.where()函數可以是按照給定的條件合並兩個tensor;
條件滿足的位置取前者,否則取后者。
- 代碼示例
>>> import torch
>>> a=torch.randn(3,5)
>>> b=torch.ones(3,5)
>>> a
tensor([[-0.0310, 1.5895, 1.6003, -1.7584, 1.1478],
[ 0.6773, 0.7763, 0.5024, 0.4952, 0.4198],
[ 1.5132, 0.5185, 0.2956, -0.6312, -1.4787]])
>>> b
tensor([[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1.]])
>>> torch.where(a>0,a,b) #合並a,b兩個tensor a>0的地方保存a的原值,否則取b的值。
tensor([[1.0000, 1.5895, 1.6003, 1.0000, 1.1478],
[0.6773, 0.7763, 0.5024, 0.4952, 0.4198],
[1.5132, 0.5185, 0.2956, 1.0000, 1.0000]])