fix device

README.md

@@ -239,19 +239,17 @@ Then, copy the code snippet below to run the example.
 
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
-device = "auto"
 model_path = "ibm-granite/granite-3.0-3b-a800m-base"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 # drop device_map if running on CPU
-model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
+model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
 model.eval()
 # change input text as desired
 input_text = "Where is the MIT-IBM Watson AI Lab located?"
 # tokenize the text
-input_tokens = tokenizer(input_text, return_tensors="pt").to(device)
+input_tokens = tokenizer(input_text, return_tensors="pt").to(model.device)
 # generate output tokens
-output = model.generate(**input_tokens,
-                        max_length=4000)
+output = model.generate(**input_tokens, max_length=4000)
 # decode output tokens into text
 output = tokenizer.batch_decode(output)
 # print output