torch.distributions.multinomial.Multinomial(total_count=1, probs=None, logits=None, validate_args=None)

the total_count is must be int, but not tensor, so I can't use this function in a batch scenario. It's quite inconvenient.

cc @vincentqb @fritzo @neerajprad @alicanb @vishwakftw