Update README.md

README.md

@@ -40,13 +40,8 @@ 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", torch_dtype=torch.float16, device_map=device_map, resume_download=True).eval()
+model = AutoModelForCausalLM.from_pretrained("m-a-p/ChatMusician", torch_dtype=torch.float16, device_map="cuda", resume_download=True).eval()
 
 generation_config = GenerationConfig(
     temperature=0.2,