log

custom_generate/generate.py

@@ -55,5 +55,6 @@ def generate(model, input_ids, generation_config=None, n_sigma:float=1.0, **kwar
         next_tokens = torch.multinomial(torch.softmax(filtered_logits, dim=-1), num_samples=1)
         input_ids = torch.cat((input_ids, next_tokens), dim=-1)
         cur_length += 1
+        print(f"Current length: {cur_length}, Next token: {next_tokens.item()}")
 
     return input_ids