Update README.md

README.md

@@ -35,11 +35,13 @@ alt="drawing" width="600"/>
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+import torch
 
 model_id = "Trendyol/Trendyol-LLM-7b-chat-v1.8"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, 
-                                             device_map='auto', 
+                                             device_map='auto',
+                                             torch_dtype= torch.bfloat16,
                                              load_in_8bit=True)
 
 sampling_params = dict(do_sample=True, temperature=0.3, top_k=50, top_p=0.9)
@@ -48,6 +50,7 @@ pipe = pipeline("text-generation",
                 model=model, 
                 tokenizer=tokenizer,
                 device_map="auto",
+                torch_dtype= torch.bfloat16,
                 max_new_tokens=1024, 
                 return_full_text=True,
                 repetition_penalty=1.1
@@ -81,6 +84,7 @@ pipe = pipeline("conversational",
                 model=model, 
                 tokenizer=tokenizer,
                 device_map="auto",
+                torch_dtype= torch.bfloat16,
                 max_new_tokens=1024,
                 repetition_penalty=1.1
                )