Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -60,7 +60,7 @@ 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)
-    probs = F.softmax(logits, dim=-1)
+    probs = F.softmax(logits.to(torch.float32), dim=-1)
     # values, indices = torch.max(probs, dim=1, keepdim=True)
     # mask = (probs == values).float()
     # probs = probs * (1 - mask)
@@ -70,8 +70,8 @@ def sample(logits, temperature: float=1.0, top_k: int=2000, top_p: float=1.0, sa
     if sample_logits:
         # add to fix 'nan' and 'inf'
         probs = torch.where(torch.isnan(probs), torch.tensor(0.0), probs)
-        probs = torch.clamp(probs, min=0, max=None)  
-        probs = probs / probs.sum()  
+        # probs = torch.clamp(probs, min=0, max=None)  
+        # probs = probs / probs.sum()  
         print(f'inf：{torch.any(torch.isinf(probs))}')
         print(f'nan: {torch.any(torch.isnan(probs))}')