Update README.md

README.md

@@ -85,6 +85,48 @@ print(tokenizer.decode(outputs[0]))
 
 </details>
 
+### Running the model on a GPU using `torch.compile`
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
+model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b", torch_dtype=torch.bfloat16).to(0)
+
+model = torch.compile(model)
+
+input_text = "Question: How many hours in one day? Answer: "
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
+<details>
+<summary> Click to expand </summary>
+
+```python
+# pip install accelerate
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-mamba-7b")
+model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b", device_map="auto")
+
+input_text = "Question: How many hours in one day? Answer: "
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids)
+print(tokenizer.decode(outputs[0]))
+```
+
+</details>
+
 ### Running the model on a GPU using different precisions
 
 #### FP16