Update app.py

app.py

@@ -19,10 +19,10 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 # device_map = infer_auto_device_map(model, dtype="float16")
 # st.write(device_map)
 
-@st.cache(allow_output_mutation=True)
+@st.cache
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True)
-    model.to(device)
+    model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True, load_in_8bit=True)
+    # model.to(device)
     tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
     return pipe