Update README.md

README.md

@@ -82,8 +82,12 @@ start = time.time()
 for i in range(200):
     next_token = model(input_ids).logits[:, -1].argmax(-1)
     generated_token_ids.append(next_token.item())
+    
+    print(next_token.item())
     input_ids = torch.cat([input_ids, next_token.unsqueeze(1)], dim=-1)
-    if "<nexa_end>" in tokenizer.decode(generated_token_ids):
+
+    # 32041 is the token id of <nexa_end>
+    if next_token.item() == 32041:
         break
 
 print(tokenizer.decode(generated_token_ids))