Search code examples
pythonpytorchtensortorch

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

Say I have a tensor A with any shape. It has a number k of non-zero elements. I want to build another tensor B, with 1s where A is non zero and 0s where A is zero.

For example:

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

then B will be :

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

Is there a simple way to implement this in Pytorch?

Solution

  • I believe it's:

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

    Also:

    B = A.bool().int()