Update app.py

app.py

@@ -16,7 +16,7 @@ st.title("Movie Reviews: An NLP Sentiment analysis")
 
 #################################################################### Cache the model loading
 
-@st.cache_data()
+@st.cache(allow_output_mutation=True)
 def load_model():
   model_pkl_file = "sentiment_model.pkl"  
   with open(model_pkl_file, 'rb') as file:  
@@ -24,12 +24,12 @@ def load_model():
   return model
 
 def load_cnn():
-  model = CNN(16236, 300, 128, [3, 8], 0.5, 2)
-  model.load_state_dict(torch.load('model_cnn.pkl'))
-  model.eval()
+    model = CNN(16236, 300, 128, [3, 8], 0.5, 2)
+    model.load_state_dict(torch.load('model_cnn.pkl', map_location=torch.device('cpu')))
+    model.eval()
+    return model
 
-  return model
-def predict_sentiment(text, model, vocab, torch_text = False):
+def predict_sentiment(text, model, vocab=16236, torch_text = False):
     tokenizer = get_tokenizer("basic_english")
     if torch_text == True:
         processor.transform(text)