Update app.py

app.py

@@ -22,7 +22,7 @@ from torch import nn
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 tokenizer = AutoTokenizer.from_pretrained("tmberooney/medllama-merged")
-model = AutoModelForCausalLM.from_pretrained("tmberooney/medllama-merged", flash_attn=True)
+model = AutoModelForCausalLM.from_pretrained("tmberooney/medllama-merged")
 '''device_map = {"transformer.word_embeddings": "cpu",
               "transformer.word_embeddings_layernorm": "cpu",
               "lm_head": "cpu",
@@ -48,7 +48,7 @@ llama_pipeline = pipeline(
     "text-generation",  # LLM task
     model=model,
     torch_dtype=torch.float16,
-    device_map="auto",
+    device_map="cpu",
     tokenizer=tokenizer
 )