快捷方式

torch.where

torch.where(condition, input, other, *, out=None) Tensor

根据 condition,返回一个从 inputother 中选择的元素的张量。

操作定义如下

outi={inputiif conditioniotheriotherwise\text{out}_i = \begin{cases} \text{input}_i & \text{if } \text{condition}_i \\ \text{other}_i & \text{otherwise} \\ \end{cases}

注意

张量 conditioninputother 必须是 可广播的

参数
  • condition (BoolTensor) – 当为 True(非零)时,生成 input,否则生成 other

  • input (TensorScalar) – 值(如果 input 是标量)或在 conditionTrue 的索引处选择的数值

  • other (张量标量) – 值(如果 other 是标量)或在 conditionFalse 的索引处选择的值

关键字参数

out (张量, 可选) – 输出张量。

返回值

一个形状等于 conditioninputother 的广播形状的张量

返回类型

张量

示例

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源