Update app.py

app.py

@@ -23,8 +23,7 @@ tokenizer = joblib.load("tokenizer.joblib") # used for LSTM
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer1 = DistilBertTokenizer.from_pretrained("tokenizer_bert")
 model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=5)
-model.load_state_dict(torch.load("fine_tuned_bert_model1.pth", map_location=torch.device('gpu')))
-model = model.to('gpu')
+model.load_state_dict(torch.load("fine_tuned_bert_model1.pth", map_location=device))
 
 # Decode label function
 # {'business': 0, 'entertainment': 1, 'health': 2, 'politics': 3, 'sport': 4}