Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -68,13 +68,13 @@ def sample(logits, temperature: float=1.0, top_k: int=2000, top_p: float=1.0, sa
     # mask = (probs == values).float()
     # probs = probs * (1 - mask)
     if sample_logits:
-        # print(f'inf：{torch.any(torch.isinf(probs))}')
-        # print(f'nan: {torch.any(torch.isnan(probs))}')
+     
         # add to fix 'nan' and 'inf'
-        probs = torch.where(torch.isnan(probs), torch.tensor(0.0), probs)
+        probs = torch.where(torch.isinf(probs), torch.tensor(0.0), probs)
         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))}')
 
         idx = torch.multinomial(probs, num_samples=1)
     else: