Does anyone know how to quickly apply a mask to a torch tensor? Say I have a tensor t_in with some values in it and a mask t_mask with values 0 and 1. In numpy I could do the following:

t_result = np.ma.array(t_in, mask=t_mask)

That would mask all the values of t_in where the value of the equivalent index of t_mask is 0. Then I can do operations such as t_result.sum() and it will only give me the sum of the non-masked values. I’m primarily looking to do just that, but cannot get my head around how it can be done efficiently with torch tensors.

If you just want to set the values to 0 when the mask is 0 then the mask and your tensor should be of the same type and just multiply the two.
If you want to extract as a new tensor all the entries of your tensor where the mask value is 1, the the mask should be a ByteTensor and you can use masked_select or []: tensor[mask] or tensor.masked_select(mask).

I recently found the need for doing masked operations on PyTorch tensors. Here’s how I implemented them (taken from my blog post):

def masked_mean(tensor, mask, dim):
"""Finding the mean along dim"""
masked = torch.mul(tensor, mask) # Apply the mask using an element-wise multiply
return masked.sum(dim=dim) / mask.sum(dim=dim) # Find the average!
def masked_max(tensor, mask, dim):
"""Finding the max along dim"""
masked = torch.mul(tensor, mask)
neg_inf = torch.zeros_like(tensor)
neg_inf[~mask] = -math.inf # Place the smallest values possible in masked positions
return (masked + neg_inf).max(dim=dim)