Masking tensor of same shape in PyTorch

Question

Given an array and mask of same shapes, I want the masked output of the same shape and containing 0 where mask is False.

For example,

# input array
img = torch.randn(2, 2)
print(img)
# tensor([[0.4684, 0.8316],
#        [0.8635, 0.4228]])
print(img.shape)
# torch.Size([2, 2])

# mask
mask = torch.BoolTensor(2, 2)
print(mask)
# tensor([[False,  True],
#        [ True,  True]])
print(mask.shape)
# torch.Size([2, 2])

# expected masked output of shape 2x2
# tensor([[0, 0.8316],
#        [0.8635, 0.4228]])

Issue: The masking changes the shape of the output as follows:

#1: shape changed
img[mask]
# tensor([0.8316, 0.8635, 0.4228])

Answer 1

Simply type-cast your boolean mask to an integer mask, followed by float to bring the mask to the same type as in img. Perform element-wise multiplication afterwards.

masked_output = img * mask.int().float()

Answer 2

One of the ways I found to solve it was:

img[mask==False] = 0

or using

img[~mask] = 0

It'll change the img itself.