Update README.md

README.md

@@ -61,9 +61,10 @@ model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-instruct-4b
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids.to("cuda")
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
 
-outputs = model.generate(input_ids)
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```
 
@@ -87,9 +88,10 @@ model = torch.compile(model)
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids.to("cuda")
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
 
-outputs = model.generate(input_ids)
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```