使用 Pytorch 中另一个张量具有非零元素的张量创建一个张量

Question

假设我有一个任意形状的张量 A。它有 k 个非零元素。我想构建另一个张量 B，A 为非零时为 1s，A 为零时为 0s。

例如：

A = [[1,2,0],
     [0,3,0],
     [0,0,5]]

那么 B 将是：

B = [[1,1,0],
     [0,1,0],
     [0,0,1]]

有没有在 Pytorch 中实现这个的简单方法？

Answer 1

我认为是：

B = (A!=0).int()

另外：

B = A.bool().int()

Create a tensor with ones where another tensor has non-zero elements in Pytorch