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 = torch.zeros(x.shape, dtype=torch.uint8)
# Your approach
for i in range(batch_size):
    mask_active[i, idx[i]] = 1

# Advanced indexing
mask_active_fast[torch.arange(batch_size).unsqueeze(1), idx] = 1

# Check
(mask_active==mask_active_fast).all()