Update README.md

README.md

@@ -36,12 +36,18 @@ You can use the models through Huggingface's Transformers library. Check our Git
 ## CLI demo
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+import torch
 from string import Template
 prompt_template = Template("Human: ${inst} </s> Assistant: ")
 
+if torch.cuda.is_available():
+    device_map = "cpu"
+else:
+    device_map = "cuda"
+
 tokenizer = AutoTokenizer.from_pretrained("m-a-p/ChatMusician", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("m-a-p/ChatMusician", trust_remote_code=True).eval()
-model.cuda()
+model = AutoModelForCausalLM.from_pretrained("m-a-p/ChatMusician", torch_dtype=torch.float16, device_map=device_map, resume_download=True).eval()
+
 generation_config = GenerationConfig(
     temperature=0.2,
     top_k=40,