Update app.py

app.py

@@ -32,7 +32,7 @@ tag_to_id, id_to_tag, id_to_description = load_tags_info()
 
 @st.cache
 def load_model():
-    return DistilBertForSequenceClassification.from_pretrained('./').to('cuda')
+    return DistilBertForSequenceClassification.from_pretrained('./')
 
 def load_tokenizer():
     return AutoTokenizer.from_pretrained(base_model_name)
@@ -68,7 +68,7 @@ with st.form("Taxonomizer"):
     if submitted:
         prompt = 'Title: ' + title + ' Abstract: ' + abstract
         tokens = tokenizer(prompt, truncation=True, padding='max_length', return_tensors='pt')['input_ids']
-        preds = softmax(model(tokens.reshape(1, -1).to('cuda')).logits / temperature, dim=1)
+        preds = softmax(model(tokens.reshape(1, -1)).logits / temperature, dim=1)
         tags = top_xx(preds, xx)
         other_tags = []
         st.header('Inferred tags:')