torch.where()

最新推荐文章于 2023-08-18 11:37:17 发布

ICDAT_GO

最新推荐文章于 2023-08-18 11:37:17 发布

阅读量708

点赞数 1

分类专栏： PyTorch学习笔记文章标签： PyTorch torch.where 条件操作张量编程

本文链接：https://blog.csdn.net/qq_40657582/article/details/121624968

版权

PyTorch学习笔记专栏收录该内容

8 篇文章 1 订阅

订阅专栏

torch.where(condition, x, y) → Tensor

Return a tensor of elements selected from either x or y, depending on condition.

condition = torch.tensor([[0.6, 0.7], [0.8, 0.4]])
print(condition)
x = torch.zeros(2,2)
print(x)
y = torch.ones(2,2)
print(y)
# 输出
tensor([[0.6000, 0.7000],
        [0.8000, 0.4000]])
tensor([[0., 0.],
        [0., 0.]])
tensor([[1., 1.],
        [1., 1.]])

print(torch.where(condition>0.5, x, y))
# 输出
tensor([[0., 0.],
        [0., 1.]])