Update app.py

app.py

@@ -9,7 +9,7 @@ dtype = torch.bfloat16
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="cpu",
+    #device_map="cpu",
     torch_dtype=dtype,
 )
 
@@ -44,7 +44,7 @@ if prompt := st.chat_input():
 
     ##Get response to the message using client
     inputs = tokenizer.encode(text, add_special_tokens=False, return_tensors="pt")
-    outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=150)
+    outputs = model.generate(**inputs, max_new_tokens=150)