Update README.md

README.md

@@ -65,7 +65,7 @@ You can run the smashed model with these steps:
    model = AutoAWQForCausalLM.from_quantized("PrunaAI/djuna-Q2.5-Veltha-14B-0.5-AWQ-4bit-smashed", trust_remote_code=True, device_map='auto')
    tokenizer = AutoTokenizer.from_pretrained("djuna/Q2.5-Veltha-14B-0.5")
     
-   input_ids = tokenizer("What is the color of prunes?,", return_tensors='pt').to(model.device)["input_ids"]
+   input_ids = tokenizer("What is the color of prunes?,", return_tensors='pt').to(next(model.parameters()).device)["input_ids"]
     
    outputs = model.generate(input_ids, max_new_tokens=216)
    tokenizer.decode(outputs[0])