Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -60,8 +60,6 @@ def sample(logits, temperature: float=1.0, top_k: int=2000, top_p: float=1.0, sa
     logits = logits[:, -1, :] / max(temperature, 1e-5)
     if top_k > 0 or top_p < 1.0:
         logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
-    print(logits.sum())
-    print(logits)
     probs = F.softmax(logits, dim=-1)
     # values, indices = torch.max(probs, dim=1, keepdim=True)
     # mask = (probs == values).float()
@@ -93,6 +91,8 @@ def logits_to_probs(logits, temperature: float = 1.0, top_p: float=1.0, top_k: i
 def prefill(model, cond_idx: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, condition:torch.Tensor, **sampling_kwargs):
     if cfg_scale > 1.0:
         logits, _ = model(None, cond_idx, input_pos, condition=condition)
+        print(logits.sum())
+        print(logits)
         logits_combined = logits
         cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0)
         logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale