self,
input,
ar_mask,
- summed_logits,
+ seq_logproba,
temperature=1.0,
deterministic_synthesis=False,
forbidden_tokens=None,
else:
dist = torch.distributions.categorical.Categorical(logits=logits)
t_next = dist.sample()
- if summed_logits is not None:
- summed_logits += logits[torch.arange(t_next.size(0)), t_next].sum(
- dim=-1
- )
+
+ all_n = torch.arange(t_next.size(0))
+ seq_logproba += logits[all_n, t_next].sum(dim=-1)
input[:, s] = ar_mask[:, s] * t_next + (1 - ar_mask[:, s]) * input[:, s]