Solved by ptrblck in post #2 
              
              
                I think you could use advanced indexing for your use case. 
Here is a small example with smaller sizes for better understanding: 
batch_size = 5
features = 4

idx = torch.empty(batch_size, 2, dtype=torch.long).random_(features)
mask_active = torch.zeros(x.shape, dtype=torch.uint8)
mask_active_fast =…