CODEWITHSUNDEEP
booleanpytorchmask
Rio1210
Rio1210

Reputation: 248

PyTorch extracting tensor elements with boolean mask (retaining dimensions)

Say, I have a PyTorch 2x2 tensor, and I also generated a boolean tensor of the same dimension (2x2). I want to use this as a mask.

For example, if I have a tensor:

tensor([[1, 3],
        [4, 7]])

And if my mask is:

tensor([[ True, False],
        [False,  True]])

I want to use that mask to get a tensor where elements corresponding to True from my original tensor are retained, whereas elements corresponding to False are set to zero in the output tensor.

Expected Output:

tensor([[1, 0],
        [0, 7]])

Any help is appreciated. Thanks!

Upvotes: 2

Views: 2784

Answers (1)

Lior Hirsch
Lior Hirsch

Reputation: 36

Assume you have :

t = torch.Tensor([[1,2], [3,4]])
mask = torch.Tensor([[True,False], [False,True]])

You can use the mask by:

masked_t = t * mask

and the output will be:

tensor([[1., 0.],
        [0., 4.]])

Upvotes: 2

Related Questions