Update app.py

app.py

@@ -5,9 +5,10 @@ import datasets
 
 @st.cache
 def load_model():
-     return AutoTokenizer.from_pretrained('distilbert-base-cased'), AutoModelForSequenceClassification.from_pretrained('./')    
+     return AutoModelForSequenceClassification.from_pretrained('./')    
 
-tokenizer, model = load_model()
+tokenizer = AutoTokenizer.from_pretrained('distilbert-base-cased')
+model = load_model()
 
 title = st.text_area('Title')
 summary = st.text_area('Summary')
@@ -15,14 +16,13 @@ summary = st.text_area('Summary')
 label_to_tag = {0: 'Computer science', 1: 'Math', 2: 'Physics',
                   3: 'Quantum biology', 4: 'Statistic'}
 
-device= 'cuda' if torch.cuda.is_available() else 'cpu'
 def predict(title, summary):
     dataset = datasets.Dataset.from_dict({'title': [title], 
                                           'summary': [summary.replace("\n", " ")]})
     dataset = tokenizer(dataset["title"], dataset['summary'],
                         padding="max_length", truncation=True, return_tensors='pt')
-    logits = model(input_ids=dataset['input_ids'].to(device),
-                   attention_mask=dataset['attention_mask'].to(device))['logits']
+    logits = model(input_ids=dataset['input_ids'],
+                   attention_mask=dataset['attention_mask'])['logits']
     probs = torch.nn.functional.softmax(logits)[0].cpu().detach()
     preds = []
     proba = 0.