if deterministic_synthesis:
t_next = logits.argmax(1)
else:
dist = torch.distributions.categorical.Categorical(logits=logits)
t_next = dist.sample()
if deterministic_synthesis:
t_next = logits.argmax(1)
else:
dist = torch.distributions.categorical.Categorical(logits=logits)
t_next = dist.sample()