torch.where¶
- torch.where(condition, input, other, *, out=None) Tensor ¶
根据
condition
,返回一个从input
或other
中选择的元素的张量。操作定义如下
注意
张量
condition
、input
、other
必须是 可广播的。- 参数
- 关键字参数
out (张量, 可选) – 输出张量。
- 返回值
一个形状等于
condition
、input
、other
的广播形状的张量- 返回类型
示例
>>> x = torch.randn(3, 2) >>> y = torch.ones(3, 2) >>> x tensor([[-0.4620, 0.3139], [ 0.3898, -0.7197], [ 0.0478, -0.1657]]) >>> torch.where(x > 0, 1.0, 0.0) tensor([[0., 1.], [1., 0.], [1., 0.]]) >>> torch.where(x > 0, x, y) tensor([[ 1.0000, 0.3139], [ 0.3898, 1.0000], [ 0.0478, 1.0000]]) >>> x = torch.randn(2, 2, dtype=torch.double) >>> x tensor([[ 1.0779, 0.0383], [-0.8785, -1.1089]], dtype=torch.float64) >>> torch.where(x > 0, x, 0.) tensor([[1.0779, 0.0383], [0.0000, 0.0000]], dtype=torch.float64)
- torch.where(condition) LongTensor 的元组
torch.where(condition)
等同于torch.nonzero(condition, as_tuple=True)
。注意
另请参阅
torch.nonzero()
。