bug: fix cache

app.py

@@ -18,13 +18,17 @@ MODELS = {
 
 headers = {}
 
-@st.cache(show_spinner=False, persist=True)
-def load_gpt(model_type, text):
-    print("Loading model...")
+@st.cache(show_spinner=False, persist=True, hash_funcs={transformers.models.gpt2.modeling_gpt2.GPT2LMHeadModel: lambda _: None})
+def load_gpt(model_type):
     model = GPT2LMHeadModel.from_pretrained(MODELS[model_type])
+
+    return model
+
+@st.cache(show_spinner=False, persist=True, hash_funcs={transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer: lambda _: None})
+def load_gpt_tokenizer(model_type):
     tokenizer = GPT2Tokenizer.from_pretrained(MODELS[model_type])
 
-    return model, tokenizer
+    return tokenizer
 
 def get_image(text: str):
     url = "https://wikisearch.uncool.ai/get_image/"
@@ -116,7 +120,8 @@ if st.button("Run"):
     with st.spinner(text="Getting results..."):
 
         st.subheader("Result")
-        model, tokenizer = load_gpt(model_name, text)
+        model = load_gpt(model_name)
+        tokenizer = load_gpt_tokenizer(model_name)
 
         input_ids = tokenizer.encode(text, return_tensors='pt')
         output = model.generate(input_ids=input_ids,