torch.where()

API和用法：

torch.where(condition, x, y) -> Tensor

condition：判断条件，一个布尔类型的张量，表示条件。

若当前index满足条件，则取x中index对应的元素

若当前index不满足条件，则取y中index对应的元素

形状：

d 是一个shape和 b，c 相同的tensor，也就是 b，c 的shape 也必须相同。

注意：a不需要和b，c一样的shape

示例：

import torch

# 创建一个布尔类型的张量，表示条件
condition = torch.tensor([True, False, True, False])

# 创建两个与 condition 形状相同的张量
x = torch.tensor([1, 2, 3, 4])
y = torch.tensor([5, 6, 7, 8])

# 使用 torch.where() 函数获取满足条件的元素索引
result = torch.where(condition, x, y)

print(result)

输出：

tensor([1, 6, 3, 8])

torch.where()函数-CSDN博客

torch.where()详解-CSDN博客

torch.where()函数解读-CSDN博客