Probably torch.multinomial would achieve a better performance for a whole batch:
batch_size = 10
weights = torch.ones(100).expand(batch_size, -1)
torch.multinomial(weights, num_samples=3, replacement=False)