Search code examples
booleanpytorchmask

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!

Solution

  • 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.]])