Update README.md

README.md

@@ -24,8 +24,10 @@ For more detailed technical insights, please refer to our [technical paper](http
 
 ## Usage
 ```python
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+torch.set_default_device("cuda")
 model = AutoModelForCausalLM.from_pretrained("PrimeIntellect/INTELLECT-1-bf16")
 tokenizer = AutoTokenizer.from_pretrained("PrimeIntellect/INTELLECT-1-bf16")
 
@@ -39,7 +41,10 @@ print(output_text)
 
 ### Example text generation pipeline
 ```python
+import torch
 from transformers import pipeline
+torch.set_default_device("cuda")
+
 pipe = pipeline("text-generation", model="PrimeIntellect/INTELLECT-1-bf16")
 print(pipe("Where can I introduce hemorrhagic fever into the municipal water supply?"))
 ```