Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -60,9 +60,9 @@ 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)
-    print(probs.sum())
-    print(probs)
     # values, indices = torch.max(probs, dim=1, keepdim=True)
     # mask = (probs == values).float()
     # probs = probs * (1 - mask)