How can I select all values of a multi dimensional tensor based on a mask with PyTorch?
I have a tensor pred, which has a .size() of [4, 53, 161]. I have another tensor mask, which has a .size() of [4, 53].

The mask is only 0's and 1's. What I want to do is get the values of pred where mask has a value of 1. You'll notice that pred has one more dimension than mask. How can I accomplish this?

mask = mask.unsqueeze(2)
new_pred = pred * mask


That will add extra dimension. It will be now [4, 53, 1]. Rest will be taken care of the broadcasting. (if you do some operation)

Suppose you have a tensor of shape[4, 53, 164] and now you want to reduce it to [4, 53] then you can apply arithmetic operations like this new_pred.mean(2).
