Update README.md

README.md

@@ -24,9 +24,10 @@ datasets:
 Get started chatting with `StableLM-Tuned-Alpha 16-bit` by using the following code snippet:
 
 ```python
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
 tokenizer = AutoTokenizer.from_pretrained("vvsotnikov/stablelm-tuned-alpha-3b-16bit")
-model = AutoModelForCausalLM.from_pretrained("vvsotnikov/stablelm-tuned-alpha-3b-16bit")
+model = AutoModelForCausalLM.from_pretrained("vvsotnikov/stablelm-tuned-alpha-3b-16bit", torch_dtype=torch.float16)
 model.cuda()
 class StopOnTokens(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool: