Pytorch 替换tensor中大于某个值的所有元素

这是在自己做多标签分类任务时用到的，最后把结果sigmoid处理，值位于【0， 1】之间，大于阈值的取1，小于阈值的取0
两种方法：

a = torch.rand((1, 3))

# 保证和a相同的维度大小
zero = torch.zeros_like(a)
one = torch.ones_like(a)

# a中大于0.5的用one(1)替换,否则a替换,即不变
a = torch.where(a > 0.5, one, a)
 
# a中小于0.5的用zero(0)替换,否则a替换,即不变
a = torch.where(a < 0.5, zero, a)

a[a > 0.5] = 1
a[a < 0.5] = 0