Update app.py

app.py

@@ -19,9 +19,9 @@ tokenizer = AutoTokenizer.from_pretrained("ua-l/gemma-2-9b-legal-steps200-merged
 model = AutoModelForCausalLM.from_pretrained(
     "ua-l/gemma-2-9b-legal-steps200-merged-16bit-uk",
     device_map='cuda',
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
 )
-compiled_model = torch.compile(model, mode="default")
+# compiled_model = torch.compile(model, mode="default")
 
 
 print('Model dtype:', model.dtype)
@@ -36,7 +36,7 @@ def predict(question):
     ### Answer:
 '''], return_tensors = "pt").to("cuda")
 
-    outputs = compiled_model.generate(**inputs, max_new_tokens = 128)
+    outputs = model.generate(**inputs, max_new_tokens = 128)
     
     results = tokenizer.batch_decode(outputs, skip_special_tokens=True)