Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -90,9 +90,10 @@ 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:
+        print(cond_idx)
+        print(input_pos)
+        print(condition)
         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