Update app.py

app.py

@@ -11,10 +11,14 @@ bert_out_address = 'model/'
 
 # Load the configuration file
 config = BertConfig.from_json_file(os.path.join(bert_out_address, "config.json"))
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # Load the pre-trained model's weights for sequence classification
 model = BertForSequenceClassification(config)
-model.load_state_dict(torch.load(os.path.join(bert_out_address, "pytorch_model.bin")))
+# model.load_state_dict(torch.load(os.path.join(bert_out_address, "pytorch_model.bin")))
+
+model.load_state_dict(torch.load(os.path.join(bert_out_address, "pytorch_model.bin"), map_location=torch.device(device)))
+
 
 # Load the tokenizer
 tokenizer = BertTokenizer.from_pretrained(bert_out_address)