Update README.md

README.md

@@ -58,7 +58,7 @@ You can run the smashed model with these steps:
     
     model_id = "PrunaAI/c4ai-command-r-plus-bnb-4bit-smashed"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto')
     
     messages = [{"role": "user", "content": "Hello, how are you?"}]
     input_ids = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")