pytorch torch.where用法

pytorch torch.where()函数用法

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

函数的作用

根据条件，返回从x,y中选择元素所组成的张量。如果满足条件，则返回x中元素。若不满足，返回y中元素。

例子