Update app.py

app.py

@@ -9,8 +9,8 @@ import time
 model_id = "google/gemma-2-27b-it"
 assistant_id = "google/gemma-2-2b-it"
 
-model = AutoModelForCausalLM.from_pretrained(model_id).to(dtype=torch.bfloat16, device="cuda")
-assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(dtype=torch.bfloat16, device="cuda")
+model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager").to(dtype=torch.bfloat16, device="cuda")
+assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id, attn_implementation="eager").to(dtype=torch.bfloat16, device="cuda")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 @spaces.GPU