Update app.py

app.py

@@ -12,18 +12,18 @@ login(token=api_key)
 # setup model
 model_id = "google/gemma-2-2b-it"
 dtype = torch.bfloat16
-
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=dtype,)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+model.to(device)
+model.eval()
 
 def poet(text):
     prompt = 'Make 25 lines, it has to be absolutely 25 lines of text no less no exception, of shakespeare based on this prompt: ' + text
     chat = [{"role": "user", "content": prompt}]
     prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+    inputs = inputs.to(device)
     outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=350)
     return tokenizer.decode(outputs[0])