Equivalent of numpy.ma.array to mask values in pytorch?

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.


Hmm ok so after digging more looks like something like this could work:

t_m_bool = t_mask.type(torch.ByteTensor)
t_result = t_in.masked_select(t_m_bool)

Would love it if someone else could confirm that this is the best approach.



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

Ok, thank you for confirming.