CODEWITHSUNDEEP
pythonpytorchtensortorch
Qubix
Qubix

Reputation: 4353

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?

Upvotes: 1

Views: 969

Answers (1)

Quang Hoang
Quang Hoang

Reputation: 150735

I believe it's:

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

Also:

B = A.bool().int()

Upvotes: 2

Related Questions