Update app.py

app.py

@@ -20,9 +20,9 @@ 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
+@st.cache(allow_output_mutation=True)
 def load_model():
-    model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True, load_in_8bit=True)
+    model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True, load_in_8bit=True, device_map='auto')
     # model.to(device)
     tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
     pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)