Update app.py

app.py

@@ -4,11 +4,12 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelWithLMHead
   
 tokenizer = AutoTokenizer.from_pretrained("gpt2-large")
+@st.cache
+def load_model(model_name):
+    model = AutoModelWithLMHead.from_pretrained("gpt2-large")
+    return model
 
-model = AutoModelWithLMHead.from_pretrained("gpt2-large")
-
-
-
+model = load_model("gpt2-large")
 
 def infer(input_ids, max_length, temperature, top_k, top_p, num_return_sequences):