I think I didn’t explain the question correctly. So let me rephrase my question. I have a set of inputs X with shape (B, H, W, C) and a probability vector p of length same as the number of channels in X. Each channel is associated with one probability. I want to keep only 50% of the channels having the high probabilities. The expected output should be of a shape (B, H, W, C/2). Is there a way to achieve this efficiently?