Anton_Z:

However, I would like to mask out only along dimension 2.


If you want to mask out only in dimension 2, then you must have a mask that matches dimension 2.
In your example, the mask must be shaped torch.Size([5000])
Only then could you do A = A[:, :, M]