Update autoregressive/models/generate.py

autoregressive/models/generate.py

@@ -71,8 +71,6 @@ def sample(logits, temperature: float=1.0, top_k: int=2000, top_p: float=1.0, sa
         # 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)   
-        print(f'inf：{torch.any(torch.isinf(probs))}')
-        print(f'nan: {torch.any(torch.isnan(probs))}')
 
         idx = torch.multinomial(probs, num_samples=1)
     else:
@@ -139,6 +137,7 @@ def decode_n_tokens(
 
 @torch.no_grad()
 def generate(model, cond, max_new_tokens, emb_masks=None, cfg_scale=1.0, cfg_interval=-1, condition=None, condition_null=None, condition_token_nums=0, **sampling_kwargs):
+    condition = condition.to(torch.float32)
     print(condition)
     if condition is not None:
         condition = model.adapter(condition)