Update app.py

app.py

@@ -19,16 +19,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # device_map = infer_auto_device_map(model, dtype="float16")
 # st.write(device_map)
 
-@st.cache
+@st.cache(allow_output_mutation=True)
 def load_model():
-    st.write('Loading model...')
     model = pipeline("text-generation", model="bryanmildort/gpt_neo_notes")
-    st.write('Success!')
     return model
    
 # model = model.to(device)
 
-
 pipe = load_model()
 
 prompt = """Admission Date:  2130-4-14     Discharge Date:  2130-4-17