Masking along one dimension

[anyfin22]

This topic has been discussed but wasn’t finished:

Masking one dimension

Suppose I have a tensor A with the following shape: torch.Size([5, 16, 5000, 3]) I also have a mask of the same shape: torch.Size([5, 16, 5000, 3]) If I apply this mask M directly to the tensor A via A = A[M] I end up with a flattened tensor with single dimension. However, I would like to mask out only along dimension 2. In other words, I would like to get a tensor of the shape torch.Size([5, 16, 5000 - N, 3]) where N is the number of entries for which mask M is False. What is the way …

I ran into a similar problem. In my case I have a tensor with [60000, 100, 400], that I need to do some operation along the third dimension. I don’t need to do any operation on most of the entries along third dimension. I was able to produce a mask of the same size that I hope reduce the tensor to something like [60000, 100, 150], but I was only able to either flatten the entire array, or keep the original dimension and make masked values zero, either way doesn’t help.

The number of unmasked values for the third dimension also varies from a few to roughly 150… for this I’m thinking making the extra values NaN would be fine.